torch.gather scatter

torch.gather

import torch
b = a.gather(dim, index)
b = torch.gather(src, 1, torch.tensor([[0,0],[1,0]]))
dim：指定轴方向，定义了填充方式。对于二维张量，dim=0表示逐列进行行填充，而dim=1表示逐列进行行填充。
* 当dim=1时，index[0][0]的元素是1，那么它想要查找a[0][1]中的元素;
* 当dim=0时，index[0][0]的元素是1，那么它想查找的a[1][0]中的元素；
index: 按照轴方向，在target张量中需要填充的位置

a = torch.arange(15).view(3, 5)
b = torch.zeros_like(a)
b[1][2] = 1
b[0][0] = 1
c = a.gather(0, b) # dim=0
d = a.gather(1, b) # dim=1

tensor([[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14]])
tensor([[1, 0, 0, 0, 0],
        [0, 0, 1, 0, 0],
        [0, 0, 0, 0, 0]])
tensor([[5, 1, 2, 3, 4],
        [0, 1, 7, 3, 4],
        [0, 1, 2, 3, 4]])
tensor([[ 1,  0,  0,  0,  0],
        [ 5,  5,  6,  5,  5],
        [10, 10, 10, 10, 10]])

ga = torch.arange(3, 12).view(3, 3)
print(ga)
index = torch.tensor([[0, 1, 2]])
tensor_1 = ga.gather(0, index)
print(tensor_1)
index = torch.tensor([[0, 1, 2]]).t()
tensor_1 = ga.gather(1, index)
print(tensor_1.t())
index = torch.tensor([[2, 1, 0]])
tensor_1 = ga.gather(0, index)
print(tensor_1)
index = torch.tensor([[0, 2],
                      [1, 2]])
tensor_2 = ga.gather(1, index)
print(tensor_2)

tensor([[ 3,  4,  5],
        [ 6,  7,  8],
        [ 9, 10, 11]])
tensor([[ 3,  7, 11]])
tensor([[ 3,  7, 11]])
tensor([[9, 7, 5]])
tensor([[3, 5],
        [7, 8]])

torch.scatter

scatter就是gather的逆过程。如果有重复放置，那么就覆盖。

当dim=1时，index[0][0]的元素是1，那么它想要查找a[0][1]中的元素;
当dim=0时，index[0][0]的元素是1，那么它想查找的a[1][0]中的元素；

a = torch.arange(10).reshape(2,5).float()
idx = torch.LongTensor([[1, 2, 1, 1, 2], [2, 0, 2, 1, 0]])
idx_1 = idx.t()#torch.tensor()
b = torch.zeros(3, 5)
b_= b.scatter(dim=0, index=idx,src=a)
b_1= b.t().scatter(dim=1, index=idx,src=a)
print(a)
print(idx)
print(b_)
print(b_1)

tensor([[0., 1., 2., 3., 4.],
        [5., 6., 7., 8., 9.]])
tensor([[1, 2, 1, 1, 2],
        [2, 0, 2, 1, 0]])
tensor([[0., 6., 0., 0., 9.],
        [0., 0., 2., 8., 0.],
        [5., 1., 7., 0., 4.]])
tensor([[0., 3., 4.],
        [9., 8., 7.],
        [0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.]])