c++ libtorch 切片 torch::indexing::Slice()的含义

rayyy92024-10-22 18:26

在LibTorch中，torch::indexing::Slice() 是一个辅助函数，用于创建一个切片对象，这个对象可以用于对Tensor进行切片操作。切片操作允许你选取Tensor中的一个子区域，而不需要复制数据。

torch::indexing::Slice() 函数可以接受几个参数，它们定义了切片的开始、结束和步长。以下是该函数的一些常见用法：

torch::indexing::Slice(start)：从索引 start 开始，直到Tensor的末尾。
torch::indexing::Slice(start, end)：从索引 start 开始，到索引 end（但不包括 end）。
torch::indexing::Slice(start, end, step)：从索引 start 开始，到索引 end（但不包括 end），以 step 为步长。

下面是使用 torch::indexing::Slice() 的一个例子：

cpp 复制代码

#include <torch/torch.h>
#include <iostream>

int main() {
    // 创建一个维度为 [5, 5] 的Tensor
    torch::Tensor tensor = torch::rand({ 5, 5 });
    std::cout << "Original Tensor:\n" << tensor << std::endl;

    // 使用 Slice() 进行切片操作
    // 选取第二行到第四行，所有列
    torch::Tensor slice1 = tensor.index({torch::indexing::Slice(1, 4)});
    std::cout << "Slice 1:\n" << slice1 << std::endl;

    // 选取所有行，第三列到第五列
    torch::Tensor slice2 = tensor.index({torch::indexing::Slice(), torch::indexing::Slice(2, 5)});
    std::cout << "Slice 2:\n" << slice2 << std::endl;

    // 选取第一行，第二列
    torch::Tensor slice3 = tensor.index({torch::indexing::Slice(0, 1), torch::indexing::Slice(1, 2)});
    std::cout << "Slice 3:\n" << slice3 << std::endl;

    return 0;
}

在上面的例子中，我们首先创建了一个5x5的Tensor，然后使用 torch::indexing::Slice() 来进行不同的切片操作：

slice1 选取了第二行到第四行，所有的列。
slice2 选取了所有的行，第三列到第五列。
slice3 选取了第一行，第二列的一个元素。

需要注意的是，切片操作不会改变原始Tensor的数据，而是返回一个新的Tensor，该Tensor与原始Tensor共享相同的数据。