函数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]是可以不同，也可以相同。大一点小一点都行