【pytorch】torch.gather用法

torch.gather(input, dim, index, *, sparse_grad=False, out=None) → Tensor

沿着由dim指定的轴，采集输入的数值

输出的数学表达式：

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

理解：
这个操作可以理解为新建一个维度和index一样的空矩阵，空矩阵中的每个值按以下方法确定：

• dim指示：采样时，input的哪一个维度是自由的。例如设dim=0，则输入tensor的第0维是自由的。
• 空矩阵中下标为[i,j,k]处的值，应在输入矩阵下标为[m,j,k]的轴上寻找。
• 其中m值为，index中下标为[i,j,k]处的值。

为方便，下面给出维度为二时的可视化图片：

当dim=0时，二维张量的行是活动的，故output[0,0]处的值应等于input[:,0]中某一行的值，index[0,0]=0，故output[0,0]=input[0,0]。其他位置的值同理。

dim = 0
input = torch.tensor([[10, 11, 12], [13, 14, 15], [16, 17, 18]])
index = torch.tensor([[0, 1, 2], [1, 2, 0]]

output = torch.gather(input, dim, index))
# tensor([[10, 14, 18],
#         [13, 17, 12]])

dim = 1
input = torch.tensor([[10, 11, 12], [13, 14, 15], [16, 17, 18]])
index = torch.tensor([[0, 1], [1, 2], [2, 0]]

output = torch.gather(input, dim, index))
# tensor([[10, 11],
#         [14, 15],
#         [18, 16]])

---

参考：

• stack overflow
• 官方文档