torch.gather

torch.gather就是能够在一个tensor中按索引取值。但是它的取值比较自由，可以第一行取0,1，第二行取1,2，第三行取2,3，只要保证每行（列）取得数量就行了。

torch.gather(input,dim,index,sparse_grad,out)

input是输入张量

dim是按行还是按哪个维度取值

index中存储取值的索引

sparse_grad代表是否需要求稀疏索引

out来存储结果。

dim=1为按行取列元素

import torch

tensor_0 = torch.arange(3, 12).view(3, 3)
print(tensor_0)
index = torch.tensor([[2, 1],[1,0]])
tensor_1 = tensor_0.gather(1, index)
print(tensor_1)

结果

tensor([[ 3,  4,  5],
        [ 6,  7,  8],
        [ 9, 10, 11]])
tensor([[5, 4],
        [7, 6]])

如果index = torch.tensor([[2, 1],[1]])就会报错，因为每个维度取得值必须相同，除非不取。

dim=0为按列取行

import torch

tensor_0 = torch.arange(3, 12).view(3, 3)
print(tensor_0)
index = torch.tensor([[2, 1],[1,0]])
tensor_1 = tensor_0.gather(0, index)
print(tensor_1)

tensor([[ 3,  4,  5],
        [ 6,  7,  8],
        [ 9, 10, 11]])
tensor([[9, 7],
        [6, 4]])

可能会有疑问为什么不是

9,6

7,4

因为是按列取得，结果也按列排。第一列取2,1，所以结果就是

9，

6。