PyTorch中gather的理解与用法

PyTorch中gather函数的理解与用法

写在前面

最近在读一个网络的源码，代码中遇到gather不理解然后百度到一位知乎https://zhuanlan.zhihu.com/p/352877584博主的解释，查看后又有了自己理解，
下面会用到上面知乎博主的一些图，如有冒犯请私信！

举个例子，更好理解

# 准备一个3X3的tensor
import torch
tensor_0 = torch.arange(3, 12).view(3, 3)
print(tensor_0)

tensor_0的内容为：

tensor([[ 3,  4,  5],
        [ 6,  7,  8],
        [ 9, 10, 11]])

torch.gather(0,index)

用行向量index，替换行索引(dim=0)

index = torch.tensor([[2, 1, 0]])
tensor_1 = tensor_0.gather(0, index)
print(tensor_1)

输出结果：

tensor([[9, 7, 5]])

理解：

对于tensor_1

[[2, 1, 0]]

“2”对应在tensor_1的坐标表示为（0,0），那么gather(0,index)就会用该数值“2”替代该数值的坐标（0,0）的行坐标“0”，得到新的坐标（2,0）在tensor_0:

[[ 3,  4,  5],
 [ 6,  7,  8],
 [ 9, 10, 11]]

里面就是“9”；同理“1”对应在tensor_1的坐标表示为（0,1），那么gather(0,index)就会用该数值“1”替代该数值的坐标（0,1）的行坐标“0”，得到新的坐标（1,1）在tensor_0中就表示为“7”；同理“0”对应在tensor_1的坐标表示为（0,2），那么gather(0,index)就会用该数值“0”替代该数值的坐标（0,2）的行坐标“0”，得到新的坐标（0,2）在tensor_0中就表示为“5”
因此输出就是[9, 7, 5]

torch.gather(1,index)

用列向量index，替换列索引(dim=1)

index = torch.tensor([[2, 1, 0]])
tensor_1 = tensor_0.gather(1, index)
print(tensor_1)

输出结果：

tensor([[5, 4, 3]])

理解：

对于tensor_1

[[2, 1, 0]]

“2”对应在tensor_1的坐标表示为（0,0），那么gather(1,index)就会用该数值“2”替代该数值的坐标（0,0）的列坐标“0”，得到新的坐标（0,2）在tensor_0:

[[ 3,  4,  5],
 [ 6,  7,  8],
 [ 9, 10, 11]]

里面就是“5”；同理“1”对应在tensor_1的坐标表示为（0,1），那么gather(1,index)就会用该数值“1”替代该数值的坐标（0,1）的列坐标“1”，得到新的坐标（0,1）在tensor_0中就表示为“4”；同理“0”对应在tensor_1的坐标表示为（0,2），那么gather(1,index)就会用该数值“0”替代该数值的坐标（0,2）的列坐标“2”，得到新的坐标（0,0）在tensor_0中就表示为“3”
因此输出就是[5, 4, 3]