python - 将张量分配给多个切片

Question

让

a = tensor([[0, 0, 0, 0],
            [0, 0, 0, 0],
            [0, 0, 0, 0]])
b = torch.tensor([1, 2])
c = tensor([[1, 2, 0, 0],
            [0, 1, 2, 0],
            [0, 0, 1, 2]])

有没有办法通过将 b 分配给 a 的切片来获得 c 而无需任何循环？也就是说，a[indices] = b 对于某些 indices 或类似的东西？

Answer 1

你可以在pytorch中使用scatter方法。

a = torch.tensor([[0, 0, 0, 0],
                 [0, 0, 0, 0],
                 [0, 0, 0, 0]])

b = torch.tensor([1, 2])

index = torch.tensor([[0,1],[1,2],[2,3]])

a.scatter_(1, index, b.view(-1,2).repeat(3,1))
# tensor([[1, 2, 0, 0],
#         [0, 1, 2, 0],
#         [0, 0, 1, 2]])