视频: https://www.bilibili.com/video/BV1my4y1u7a8?from=search&seid=11271863612450212285
官方文档: https://pytorch.org/docs/stable/generated/torch.gather.html
https://zhuanlan.zhihu.com/p/352877584
why gather
?
使用索引只能取出连续的值和单个值,不能取出不连续的多个值。
How gather
work?
使用知乎上的一个例子:图解PyTorch中的torch.gather函数
import torch
tensor_0 = torch.arange(3, 12).view(3, 3)
print(tensor_0)
输出结果为
tensor([[ 3, 4, 5],
[ 6, 7, 8],
[ 9, 10, 11]])
现在想获取对角线上的数,也就是[9,5,7]
,没有办法通过索引的方式一步进行。
如果想使用gather
要先定义index
,index
需要和tensor_0
是相同维度,也就是都是2维的。
然后通过下面的代码就可以取出对角元素。
index = torch.tensor([[2, 1, 0]])
tensor_1 = tensor_0.gather(0, index)
print(tensor_1)
输出结果
tensor([[9, 7, 5]])
图解
在代码中我们定义了index
,先为index
生成下标,也就是图中的index of index
:(0,0) (0,1) (0,2)。
由于gather
函数中设置的维度是0,所以用index
的数代替index of index
对应的数gather index
:(2,0) (1,1) (0,2)
然后根据生成的索引进行取数
官方文档
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
注意out
和index
是相同size的