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

问题简化实现torch.gather用法TIPS逆操作torch.Tensor.scatter

问题简化

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维向量

逆操作torch.Tensor.scatter

这个用法是将src的元素按照index赋给当前tensor，官方文档例子很清楚。 这里说明一下它在onehot向量生成上的应用。

y_onehot = torch.FloatTensor(3, 3) y_onehot.zero_() index = torch.arange(y_onehot.size(0)).unsqueeze(1) y_onehot.scatter_(1, index, 1.0)