pytorch 笔记：gather 函数

Posted 2022-02-02 UQI-LIUWJ

torch.gather(input, dim, index, out=None) → Tensor

我们直接用例子来说明好了

以三位张量为例

out[i] [j] [k] = tensor[index[i][j][k]]  [j]  [k] 
# dim=0

out[i] [j] [k] = tensor[i]  [index[i][j][k]]  [k] 
# dim=1

out[i] [j] [k] = tensor[i]  [j]  [index[i][j][k]] 
# dim=3

import torch
a = torch.Tensor([[1,2],[3,4]])
torch.gather(a,
            0,
            index=torch.LongTensor([[0,0],[1,0]]))
'''
tensor([[1., 2.],
        [3., 2.]])
'''

这个怎么看呢

out[0][0]	a[index[0][0]]  [0]]	a[0][0]=1
out[1][0]	a[index[1][0]]  [0]]	a[1][0] =3
out[0][1]	a[index[0][1]]  [1]]	a[0][1]=2
out[1][1]	a[index[1][1]]  [1]]	a[0][1]=2

还有一种用法在每一行选择一个的时候比较常用

import torch
a = torch.Tensor([[1,2,5,6],[3,4,7,8]])
torch.gather(a,1,torch.LongTensor([[0],[3]]))
'''
tensor([[1.],
        [8.]])
'''