Question

在pytorch中，我有一个多维张量，称其为X

X = [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], ...]

现在我想像这样为每行选择一个不同的列索引

indices = [[0], [1], [0], [2], ...]
# now I expect following values to be returned:
[[1], [5], [7], [12], ...]

我也想实现相反的效果，以便对于给定的索引我得到

[[2, 3], [4, 6], [8, 9], [10, 11]]

有没有“无”循环的“简单”方法来实现这一目标？我将不胜感激。

Answer 1

使用numpy可以轻松完成，如果需要张量，它可以来回转换

X = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
X2= X[np.arange(X.shape[0]),[0,1,0,2]]
X3 = np.setdiff1d(X, X2).reshape(X.shape[0],X.shape[1]-1)

编辑

在gpu中做同样的事情，使用张量

import torch.tensor as tensor
import numpy as np

X = tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
X2= X[np.arange(X.shape[0]),[0,1,0,2]]

def allothers(X, num):
    return [x for x in range(X.shape[1]) if x not in [num]]


X3 = X[ [[x] for x in np.arange(X.shape[0])], [allothers(X, 0),allothers(X, 2),allothers(X, 1),allothers(X, 1)] ]

Answer 2

实际上，torch.gather函数执行此操作。

例如

a = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
indices = torch.tensor([[0], [1], [0], [2]])
a.gather(1, indices)

将完全返回

tensor([[ 1],
        [ 5],
        [ 7],
        [12]])

我不再需要相反的内容，为此，我建议只创建一个全为掩码的遮罩，然后将“聚集”张量的各个索引设置为0或仅创建一个包含各自的“聚集”张量的新索引。相反的键。例如：

indices_opposite = [np.setdiff1d(np.arange(a.size(1)), i) for i in indices.numpy()]

在每行中选择/屏蔽不同的列索引

2 个答案: