Torch tensor.gather()学习

最近在学习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，也就是行了。

大概就是可以这样理解吧，本人也是初次学习，希望大家能多多交流~