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