pytorch的gather函数个人理解

首先写一下官方的案例

b = torch.Tensor([[1,2,3],[4,5,6]])
print b
index_1 = torch.LongTensor([[0,1],[2,0]])
index_2 = torch.LongTensor([[0,1,1],[0,0,0]])
print torch.gather(b, dim=1, index=index_1)
print torch.gather(b, dim=0, index=index_2)

这是输出结果：

 1  2  3
 4  5  6
[torch.FloatTensor of size 2x3]


 1  2
 6  4
[torch.FloatTensor of size 2x2]


 1  5  6
 1  2  3
[torch.FloatTensor of size 2x3]

这是官方文档的解释：

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]

就个人理解的话，dim=1表示索引是列，那么对于index_1=[[0,1],[2,0]].那么对于gather这个函数就是，index_1每一个元素的行对应b里面的行(dim=1表示不改变行)，然后里面的每一个元素的取值对应b里面的列。
也就是说：
0 一>0列0行一>1
1 一>1列0行一>2
2 一>2列1行一>6
1 一>1列1行一>4
所以输出结果是[[1,2][6,4]]

dim=0表示索引是行，那么对于index_2=[[0,1,1],[0,0,0]]。那么对于gather这个函数就是，index_2的每一个元素的列对应b里面的列，然后里面的每一个元素的取值对应b里面的行。
也就是说：
0 一>0行0列一>1
1 一>1行1列一>5
1 一>1行2列一>6
0 一>0行0列一>1
0 一>0行1列一>2
0 一>0行2列一>3
所以输出结果是[[1,5,6][1,2,3]]