Torch.gather的理解

Torch.gather方法通常运用在分类问题中，最后分类结果为onehot格式时，对于真实标签的提取。官方文档的公式理解有点难度。以下简单做一下思考。

代码引用自 Pytorch中的torch.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]

gather是一个能够按照指定维度和方法进行索引的工具。

在torch.gather(b, dim=1, index=index_1)中，第一个参数代表待索引的原始向量，也就是说在输出中的所有数都出自这个向量中。

第二个dim=0/1/2指的是按照第几维进行索引。

index参数指的是如何进行索引。index向量中的每一个数都是代表该数位置的原向量的值，所以最终输出的结果是和index向量保持一致的。

先看torch.gather(b, dim=1, index=index_1)的输出。是从第1维（开始下标为0）进行索引。第一维是列，第二维是行。也就是说这里面的数字分别代表列号。index数组中的每一个数都代表着第1维位置的数。直接举例子来说：

第一行 [0,1] 指的是这两个数需要从原数组中的第一行的对应位置来取。0是第一行的第一个，为1；1是第一行的第2个，为2；以此类推，第二行2，0，需要去原数组第二行找，即可找到第二行列号为2的元素是6，列号为0的是4。

再看dim=0的情况：dim=0代表index数组中的每一个数都代表着第0维位置的数。 index数组第一行第一个0，去原来的数组的第一列，找到行号为0的数，就是1；第一行第二个为1，去看原数组第二列，找到行号为1的也就是5；后面的1，找第三列，行号为1的数，就是6；下面的三个0，分别去找1，2，3列行号为0的数，就是1，2，3. 所以有了上文中的输出。

总结：index中的数都是代表着某一维特定位置的元素。dim=0时，按照列来看；dim=1时按照行来看。在应用中，index设置为按照最大值提取，就可以把softmax输出最终的类别给提取出来。