Question

我正在使用C++ frontend for PyTorch，并且遇到了相对基本的索引问题。

我有一个8的{{1}}张量，例如下面的张量：

我还有另一个Tensor，其中仅包含[ Variable[CUDAFloatType]{8,6} ] 0 1 2 3 4 5 0 1.7107e-14 4.0448e-17 4.9708e-06 1.1664e-08 9.9999e-01 2.1857e-20 1 1.8288e-14 5.9356e-17 5.3042e-06 1.2369e-08 9.9999e-01 2.4799e-20 2 2.6828e-04 9.0390e-18 1.7517e-02 1.0529e-03 9.8116e-01 6.7854e-26 3 5.7521e-10 3.1037e-11 1.5021e-03 1.2304e-06 9.9850e-01 1.4888e-17 4 1.7811e-13 1.8383e-15 1.6733e-05 3.8466e-08 9.9998e-01 5.2815e-20 5 9.6191e-06 2.6217e-23 3.1345e-02 2.3024e-04 9.6842e-01 2.9435e-34 6 2.2653e-04 8.4642e-18 1.6085e-02 9.7405e-04 9.8271e-01 6.3059e-26 7 3.8951e-14 2.9903e-16 8.3518e-06 1.7974e-08 9.9999e-01 3.6993e-20个元素，例如：

我想用第二个张量来索引第一个张量的行：

[ Variable[CUDALongType]{8} ] 
 0
 3
 4
 4
 4
 4
 4
 4

我尝试了几种不同的方法，包括index_select，但似乎产生的输出与输入（0 0 1.7107e-14 1 1.2369e-08 2 9.8116e-01 3 9.9850e-01 4 9.9998e-01 5 9.6842e-01 6 9.8271e-01 7 9.9999e-01）的尺寸相同。

在Python中，我认为我可以使用Python的内置索引编制索引，如下所述：https://github.com/pytorch/pytorch/issues/1080

不幸的是，在C ++中，我只能用标量（零维Tensor）索引一个Tensor，所以我认为这种方法不适用于我。

如何在不求助于循环的情况下达到期望的结果？

Answer 1

事实证明，您可以通过几种不同的方式来执行此操作。一个带有gather，另一个带有index。在PyTorch discussions中，我问了同样的问题：

使用torch::gather

auto x = torch::randn({8, 6});
int64_t idx_data[8] = { 0, 3, 4, 4, 4, 4, 4, 4 };
auto idx = x.type().toScalarType(torch::kLong).tensorFromBlob(idx_data, 8);
auto result = x.gather(1, idx.unsqueeze(1));

使用特定于C ++的torch::index

auto x = torch::randn({8, 6});
int64_t idx_data[8] = { 0, 3, 4, 4, 4, 4, 4, 4 };
auto idx = x.type().toScalarType(torch::kLong).tensorFromBlob(idx_data, 8);
auto rows = torch::arange(0, x.size(0), torch::kLong);
auto result = x.index({rows, idx});

PyTorch for C ++中的按行元素索引

1 个答案: