本文主要是介绍pytorch中gather函数的理解,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!
pytorch函数gather理解
torch.gather(input, dim, index, out=None) → Tensor
Parameters:
- input (Tensor) – 源张量
- dim (int) – 索引的轴
- index (LongTensor) – 聚合元素的下标(index需要是torch.longTensor类型)
- out (Tensor, optional) – 目标张量
公式含义
这个函数的意义就是可以重新排列特定维度的信息。对一个三维张量,从公式来看,输出是下面这种,就是在特定维度上,用索引index下标代替所在位置的值。
out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0
out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1
out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2
直观理解
原始tensor ,名称为a
a = torch.randint(0, 30, (2, 3, 5))
以下以 CxHxW的维度讲述,其中C=2,H=3, W=5,
index = torch.LongTensor([[[0,1,2,0,2],[0,0,0,0,0],[1,1,1,1,1]],[[1,2,2,2,2],[0,0,0,0,0],[2,2,2,2,2]]])
指定dim = 1,也就是在第二个维度上H重新排列,
b = torch.gather(a, 1,index)
此时,第一个维度C是不会改变的,还是存在两个通道C,分别是a[0]和a[1],
针对a[0]或者a[1] , 在高度维度H上,分别是3行,a[0][0:2] a[1][0:2]。即
a[0].shape == [3,5]
因此,如果选择dim=1,则index 张量里面的数必须在0-2之间,不然会越界,
下一步就是选取数字了。
针对每一个通道C,输出张量b,只需要按照index重新排列矩阵即可
例如在第b[0,1,2]的位置,则选择a[0][index[0,1,2]][2]的值进行代替即可。
同理在其他维度也是一样。
注意点
需要注意的是索引矩阵不能越界,例如针对上述a[2,3,5],
如果指定dim=0,则index里面的数不能超过1,指定dim=1,则index不能超过2,指定dim=3,则index不能超过4
本文参考https://www.jianshu.com/p/5d1f8cd5fe31
这篇关于pytorch中gather函数的理解的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!