其他分享
首页 > 其他分享> > Torch.gather的理解

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时按照行来看。所以在传参时也要注意。当设置dim=0时,按照列来看。在应用中,index设置为按照最大值提取,就可以把softmax输出最终的类别给提取出来。

标签:index,torch,数组,dim,Torch,gather,理解,第一行
来源: https://blog.csdn.net/weixin_43537073/article/details/121292709