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

1  2
6  4

1  5  6
1  2  3

仔细研究一下结果是怎么来的, 对应原数组的索引位置   

axis = 1

[  1,  2              ->        [ [0, 0], [0, 1] 

   6,  4 ]            ->          [1, 2], [1, 0] ]

结果就是在第0行中取第0, 1元素，在第1行取第2, 0元素

可以发现  axis = 1 结果  就是 自动添加行索引，在每一行中取值 

 

再看一下axis=0时

[1, 5, 6,           ->        [ [0, 0], [1, 1], [1, 2]  

 1, 2, 3]           ->         [0, 0], [0, 1], [0, 2] ]

在第0列中取第0, 0 元素，在第一列中取1, 0 元素，第二列中取1, 0 元素

可以发现 axis = 0 就是自动添加列索引，在每一列中取值