Pytorch:gather函数

TORCH.GATHER
作用：从张量中根据下标位置沿指定dim取出指定元素聚合成新张量

#函数定义
torch.gather(input, dim, index, out=None, sparse_grad=False) → Tensor

注意蓝绿色框内的下标，也就可以明白该函数的使用。
由红框中的内容可知，out的大小和index的大小是一样的。index需要和input具有除了dim以外相同的维度数。

Example:

#一维
t = torch.tensor([1, 3, 5, 0, 9])
need_elem = torch.gather(t, 0, torch.tensor([0,4,1]))
print(need_elem)
>>>tensor([1, 9, 3])

#二维
#除了dim之外的维度不相同，报错
t = torch.tensor([[1,2],[3,4],[9,5]]) #[3,2]
need_elem = torch.gather(t, 1, torch.tensor([[0,0],[1,0]]))#index=[2,2]
print(need_elem)
>>> RuntimeError: Expected tensor [2, 2], src [3, 2] and index [2, 2] to have the same size apart from dimension 1
#正确
t = torch.tensor([[1,2],[3,4],[9,5]])#[3,2]
need_elem = torch.gather(t, 1, torch.tensor([[0],[1],[1]]))#index=[3,1]
print(need_elem)#[3,1]