最近在学习Pytorch的时候遇见了一个很棘手的函数gather,之前也转载了一个博客,但是也似懂非懂的,今天在此总结下,以便后续需要~
首先贴一下官方的介绍:
torch.gather(input, dim, index, out=None) → Tensor
Gathers values along an axis specified by dim.(沿着dim指定的轴聚集 s值。)
For a 3-D tensor the output is specified by:
out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0
out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1
out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2
Example:
t = torch.tensor([[1,2],[3,4]])
torch.gather(t, 1, torch.tensor([[0,0],[1,0]]))
tensor([[ 1, 1],
[ 4, 3]])
以上就是官方的介绍,下面进行下理解吧~
[Tensor] gather (dim,indices) – 首先这个函数需要重新分配内存。
– 该函数的功能主要是沿着dim维度,在每一个row上按照indices选取数值,indices为LongTensor类型
听起来是不是很抽象,看看下面这个图示就很好理解了:
左图是:result=src.gather(1,index), index=torch.LongTensor({{1,2,3},{2,3,1}})
其输出result的大小和index相同,result[{1,{}}]
为src每一列上的index[{{1},{}}
]对应的元素,为什么是列呢,因为dim=1,决定了沿着row数,也就是列了。
右图是:result=src.gather(2,index), index=torch.LongTensor({{1,2,3},{2,3,1}})
其输出result的大小和index相同,result[{{},1}]
为src每一行上的index[{{},{1}}
]对应的元素,为什么是行呢,因为dim=2,决定了沿着Column,也就是行了。
大概就是可以这样理解吧,本人也是初次学习,希望大家能多多交流~