PyTorch中gather的理解

2维tensor的gather操作理解
考虑2维tensor a

a = t.arange(0, 16).view(4, 4)
print(a)

output:

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

index与[[2,2,0,1]]的维度相同，即(1,4).
dim=0，按列查看，第一个2处于该张量的第0列 ，所以第一个元素取a中第0列的索引为2的元素8；
dim=1，按行查看，第一个2处于第0行，所以第一个元素取a中第0行的索引为2的元素， 由于index只有第0行，所以对应的返回为a中第0行的索引为2，2，0，1的元素，即[[2,2,0,1]].

index = t.LongTensor([[2,2,0,1]])
print('dim=0:',a.gather(0, index))
print('dim=1:',a.gather(1, index))

output:

dim=0: tensor([[8, 9, 2, 7]])
dim=1: tensor([[2, 2, 0, 1]])

dim=0，数字i在第x列就选取原张量的第x列索引为i的元素.
举例， 0表示第0列索引为0(0),3表示第1列索引为3(13)
dim=1，数字在第x行就选取原张量的第x行索引为i的元素.
举例， 0表示第0行索引为0的数(0),3表示第2行索引为3的元素(11)

index = t.LongTensor([[0,2],
                      [1,2],
                      [1,3]])
print('dim=0:\n',a.gather(0, index))
print('dim=1:\n',a.gather(1, index))

output:

dim=0:
 tensor([[ 0,  9],
        [ 4,  9],
        [ 4, 13]])
dim=1:
 tensor([[ 0,  2],
        [ 5,  6],
        [ 9, 11]])

自用笔记，记录备查