libtorch/Pytorch：用 C++ 从 Tensor 中切片通道

Question

我有一个用于分割任务的神经网络，它产生形状为 [1, 2, 416, 416] 的输出张量。第一个索引是批量大小，第二个索引是输出通道。我可以通过使用squeeze()来消除批量维度。我可以通过在squeeze()之后使用argmax(0)将两个通道减少到一个，但我想知道如何分别获取通道。

这是我当前用于将张量转换为 cv::mat 的代码片段：

tensor = tensor.squeeze().argmax(0).detach(); //instead of argmax get specific channel.
std::cout << tensor.sizes() << std::endl;
tensor = tensor.contiguous();
tensor = tensor.mul(255).clamp(0, 255).to(torch::kU8);
tensor = tensor.to(torch::kCPU);
int64_t height = tensor.size(0);
int64_t width = tensor.size(1);
cv::Mat mat = cv::Mat(cv::Size(width, height), CV_8UC1, tensor.data_ptr());

我目前使用 PyTorch/libtorch 1.2，因为它是与 VS2015 兼容的最新版本。

Answer 1

您可以使用

torch::index

:

tensor = tensor.squeeze(0).index({0});

或

torch::narrow

适用于更一般的情况，即您想要从该维度提取给定数量的行：

tensor = tensor.squeeze(0).narrow(/*dim*/=0, /*start*/=0, /*length=*/1);