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

在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]]

---

有没有一种“简单”的方法可以在没有for循环的情况下实现这一点？如果您有任何想法，我将不胜感激。

使用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)] ]

事实上，该函数正好执行此操作

比如说

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]])

我不再需要相反的，但为此，我建议只创建一个包含所有1的掩码，然后将“聚集”张量的相应索引设置为0，或者只创建一个包含相应相反键的新“聚集”张量。例如：

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

---

谢谢，但我的问题是，我在一个GPU上，想把我的数据保存在那里。