pytorch的gather()方法详解

首先，先将结果展示出来，后续根据结果来进行分析

t = torch.Tensor([[1,2,3],[4,5,6]])
index_a = torch.LongTensor([[0,0],[0,1]])
index_b = torch.LongTensor([[0,1,1],[1,0,0]])
print(t)
print(torch.gather(t,dim=1,index=index_a))
print(torch.gather(t,dim=0,index=index_b))

运行结果如下：

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

结果分析

首先，三个参数介绍
输入的变量input、指定在某一维上聚合的dim、聚合使用的索引index，输出为Tensor类型的结果（index必须为LongTensor类型）

1.当dim = 1 时，是在第二个维度进行融合，index是所需值在原矩阵中的索引。即按照列来进行融合，第几行列号为[x,x]的值。[0, 0]表示第一行第1列的值，两个都为1.[0, 1]表示对应原数组第二行第一列和第二列的值，分别为4，5.
因此最终结果如下：

tensor([[1., 1.],
        [4., 5.]])

2. 当dim = 0时，是在第一个维度进行融合，按行来进行融合。第几列行号为[x,x,x]的值。[0,1,1]表示对应原数组第一列第一行（0）的值1、第二列第二行（1）的值5，第三列第二行（1）的值6.
   最终取值如下

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