函数torch.gather()的实例记载学习

记录一下

因为总是使用这个函数，而且贼恶心，总忘记。

二维案例

1.1 二维简单案例

# 1.
x = torch.randint(0,20,(4,5)) 
x, x.shape
'''
(tensor([[ 6, 19,  3,  4,  5],
         [12,  2,  9, 13,  2],
         [ 8, 18, 16, 14,  6],
         [ 8,  8, 10, 12,  8]]),
 torch.Size([4, 5]))
'''
# 2.
index = torch.tensor([[1,1,1,1],[2,2,2,2],[3,3,3,3],[0,0,0,0]])
index,index.shape
'''
(tensor([[1, 1, 1, 1],
         [2, 2, 2, 2],
         [3, 3, 3, 3],
         [0, 0, 0, 0]]),
----------------------------------------------------
----> 注意index的 size() 是 torch.Size([4, 4])) <----
----------------------------------------------------
'''
# 3.
torch.gather(x, 1, index)
'''
tensor([[19, 19, 19, 19],
        [ 9,  9,  9,  9],
        [14, 14, 14, 14],
        [ 8,  8,  8,  8]])
'''
# 4.很明显
----------------------------------------------------
'返回 ---> size=(4,4) <--- 的矩阵': 和上面index的尺寸一样
----------------------------------------------------
[[ x[0,1]=19, x[0,1]=19, x[0,1]=19, x[0,1]=19],
 [ x[1,2]=9, x[1,2]=9, x[1,2]=9, x[1,2]=9],
 [...],
 [...]
]

1.2 二维复杂案例

# 1.
'''
(tensor([[ 6, 19,  3,  4,  5],
         [12,  2,  9, 13,  2],
         [ 8, 18, 16, 14,  6],
         [ 8,  8, 10, 12,  8]]),
 torch.Size([4, 5]))
'''
# 2.
index = torch.tensor([[1,2,3,0],[2,1,3,0],[3,3,3,3],[0,0,0,0]])
index,index.shape
'''
(tensor([[1, 2, 3, 0],
         [2, 1, 3, 0],
         [3, 3, 3, 3],
         [0, 0, 0, 0]]),
----------------------------------------------------
----> 注意index的 size() 是 torch.Size([4, 4])) <----
----------------------------------------------------
'''
# 3.
torch.gather(x, 1, index)
'''
tensor([[19,  3,  4,  6],
        [ 9,  2, 13, 12],
        [14, 14, 14, 14],
        [ 8,  8,  8,  8]])
'''
# 4.很明显
----------------------------------------------------
'返回 ---> size=(4,4) <--- 的矩阵': 和上面index的尺寸一样,这个绝对不会变
----------------------------------------------------
[[ x[0,1]=19, x[0,2]=3, x[0,3]=4, x[0,0]=6],
 [ ... ],
 [ ... ],
 [ ... ]
]
# 5. 那么当dim=0呢
'RuntimeError: Size does not match at dimension 1 get 5 vs 4'
会报错
这是因为 index.size()[1] 是 4， 而 x.size()[1] 是 5
但是选定的·dim=0·,也就是 index.size()[dim] 和 x.size()[0]是可以不同，也可以相同。大一点小一点都行