带张量的 Pytorch 索引

Question

我有一个二维张量 arr，所有条目均为 0。我有第二个张量 idx。我想将 arr 中的所有条目与 idx 中的索引合并为 1.

arr = torch.zeros(size = (2,10))
idx = torch.Tensor([
    [0,2],
    [4,5]
])
arr[idx] = 1 #This doesn't  work
print(arr)

输出应如下所示：

tensor([[1., 0., 1., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 1., 1., 0., 0., 0., 0.]])

我很有信心我肯定会在 SO 中找到其他人问这个问题，但是我找不到。我希望它不是重复的。

Answer 1

沿 dim=1 使用 scatter() 或在这种情况下最里面的维度，即 dim=-1。请注意，代替 src 张量，我只是传递了常数值 1.

In [31]: arr = torch.zeros(size=(2, 10))
 
In [32]: idx = torch.tensor([
     ...:     [0, 2],
     ...:     [4, 5]
     ...:     ])
 
In [33]: torch.scatter(arr, 1, idx, 1)
Out[33]: 
 tensor([[1., 0., 1., 0., 0., 0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 1., 1., 0., 0., 0., 0.]])

In [34]: torch.scatter(arr, -1, idx, 1)
Out[34]: 
 tensor([[1., 0., 1., 0., 0., 0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 1., 1., 0., 0., 0., 0.]])

带张量的 Pytorch 索引

Pytorch index with Tensor

python

pytorch