使用torch.gather选择tensor某一维度指定的元素

问题简化

PyTorch中，对于某个二维tensor，要求取出其每一行的指定元素，组成新的tensor。例如，对于

A = torch.tensor([[1,2,3],[4,5,6],[7,8,9]])
# tensor([[1, 2, 3],
          [4, 5, 6],
          [7, 8, 9]])

要求取出A的第0行的第1个、第1行的第2个、第2行的第0个元素。

实现

idx = torch.LongTensor([[1],[2],[0]])
out = torch.gather(A, 1, idx)
# tensor([[2],
          [6],
          [7]])

torch.gather用法

out = torch.gather(input, dim, index, out=None)
# 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
# ......

主要作用是精准地获取tensor中的元素。
注意，参数index在shape上，除了第dim维的长度可以与input不一样之外，其他维度的长度都应该保持一致。输出的shape与index一致。

TIPS

还有一种不太优雅很原始但还比较灵活的索引方式。

xt = torch.rand(3, 5, 8)
xt[[0,1,2],[1,2,4]]	# 表示分别取前3行的第1、2、4个8维向量