Pythorch gather函数

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

    Gathers values along an axis specified by dim.

    For a 3-D tensor the output is specified by:

    out[i][j][k] = input[index[i][j][k]][j][k]  # dim=0
    out[i][j][k] = input[i][index[i][j][k]][k]  # dim=1
    out[i][j][k] = input[i][j][index[i][j][k]]  # dim=2

    Parameters:	

        input (Tensor) – The source tensor
        dim (int) – The axis along which to index
        index (LongTensor) – The indices of elements to gather
        out (Tensor, optional) – Destination tensor

    Example:

    >>> t = torch.Tensor([[1,2],[3,4]])
    >>> torch.gather(t, 1, torch.LongTensor([[0,0],[1,0]]))
     1  1
     4  3
    [torch.FloatTensor of size 2x2]

由上面的公式可以看出，输出的维度和index的维度是一致的。
对于dim=0时，out[i][j][k]的取值过程为：
首先找到index[i][j][k]的值，然后作为input的第一个维度的索引，后面维度的索引j,k。最后按照索引找到input对应位置的值，作为out[i][j][k]的值。
对于dim=1时，首先找到index[i][j][k]的值，然后作为input的第二个维度的索引，其他维度的索引为i,k。最后按照索引找到input对应位置的值，作为out[i][j][k]的值。
对于dim=2时，同理。

示例

import  torch

b = torch.Tensor([[1,2,3],[4,5,6]])

print(b)
index_2 = torch.LongTensor([[0,1,1],[0,0,0]])

a = torch.gather(b,dim=0,index=index_2)
print(a)

输出为

tensor([[1., 2., 3.],
        [4., 5., 6.]])
tensor([[1., 5., 6.],
        [1., 2., 3.]])

此时dim=0，是将index对应位置的值，作为input第一个维度的索引。比如，index第一行第二列[0][1]的值为1，那么out的第一行第二列的值为inpu第二行第二列的值5.