官方示例
链接: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 中的值作为索引传递,指定从每个“行”中获取的值。 输出张量的维度与索引张量的维度相同。 以下图片能更清楚地解释它: