一看就懂的torch.gather()用法

官方文档中对此函数解释如下：
https://pytorch.org/docs/stable/generated/torch.gather.html?highlight=gather#torch.gather
其中Gathers values along an axis specified by dim.指出了gather函数本质是对一个张量进行取出其中元素的效果，类比于二维数组中通过index下标取出二维数组的元素，例如arr[1][2]=5

官方文档给出了计算式子，这里是实际计算的一个例子。
input：
[[1, 2],
[3, 4]]

index：
[[0, 0],
[1, 0]]

out的计算式如下：
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

torch.gather(input, 1, index)

out[i][j]=input[i][index[i][j]]
out[0][0]=input[0][index[0][0]]=input[0][0]=1
out[0][1]=input[0][index[0][1]]=input[0][0]=1
out[1][0]=input[1][index[1][0]]=input[1][1]=4
out[1][1]=input[1][index[1][1]]=input[1][0]=3

综上，该函数作用为Gathers values along an axis specified by dim，翻译过来也许你会看着迷糊。本质就是从完整数据中按索引取值，和二维数组一样

参考：https://zhuanlan.zhihu.com/p/352877584