PyTorch for C++ 中的逐行元素索引答案

【问题标题】：Row-wise Element Indexing in PyTorch for C++PyTorch for C++ 中的逐行元素索引
【发布时间】：2018-11-29 13:16:16
【问题描述】：

我正在使用C++ frontend for PyTorch 并且正在努力解决一个相对基本的索引问题。

我有一个8 by 6 张量，如下所示：

[ 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

我有另一个张量，其中只有 8 元素，例如：

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

我想使用第二个张量来索引我的第一个张量的行以生成：

        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

我尝试了几种不同的方法，包括 index_select，但它似乎生成的输出与输入 (8x6) 具有相同的尺寸。

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

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

如何在不使用循环的情况下达到我想要的结果？

【问题讨论】：

你想看pytorch.org/docs/stable/torch.html#torch.gather而不是index_select

标签： c++ pytorch

【解决方案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});

【讨论】：