【Torch API】torch.gather()用法详解

官方示例

链接：https://pytorch.org/docs/stable/generated/torch.gather.html

torch.gather(input, dim, index, *, sparse_grad=False, out=None) → Tensor
# 沿由 dim 指定的轴收集input的值，其输出形状与index相同。

• input (Tensor) – the source tensor
• dim (int) – the axis along which to index
• index (LongTensor) – the indices of elements to gather

对于一个3维的tensor，其输出如下：

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

图解

参考：https://stackoverflow.com/questions/50999977/what-does-the-gather-function-do-in-pytorch-in-layman-terms

torch.gather 通过沿输入维度 dim 从每一行获取值，从输入张量创建一个新张量。 torch.LongTensor 中的值作为索引传递，指定从每个“行”中获取的值。 输出张量的维度与索引张量的维度相同。 以下图片能更清楚地解释它：