我正在使用C++ frontend for PyTorch并且正在努力解决相对基本的索引问题 .
我有一个 8
by 6
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
我有另一个Tensor,里面只有 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中我只能用标量(零维张量)索引Tensor,所以我不认为这种方法对我有用 .
如何在不诉诸循环的情况下实现理想的结果?
1 回答
事实证明,你可以通过几种不同的方式做到这一点 . 一个是
gather
,另一个是index
. 从PyTorch discussions我问同样的问题:使用torch::gather
使用C特定的torch::index