在 pytorch 张量中，return 特定值行的索引数组

Question

给定以下张量，其中包含全零向量和包含 1 和 0 的向量：

tensor([[0., 0., 0., 0.],
    [0., 1., 1., 0.],
    [0., 0., 0., 0.],
    [0., 0., 1., 0.],
    [0., 0., 0., 0.],
    [0., 0., 1., 0.],
    [1., 0., 0., 1.],
    [0., 0., 0., 0.],...])

我怎样才能得到包含 1 和 0 的向量索引数组，以便输出如下所示：

indices = tensor([ 1, 3, 5, 6,...])

更新

一种方法是：

indices = torch.unique(torch.nonzero(y>0,as_tuple=True)[0])

但我不确定是否有更好的方法。

Answer 1

另一种方法是使用 torch.Tensor.any coupled with torch.Tensor.nonzero:

>>> x.any(1).nonzero()[:,0]
tensor([1, 3, 5, 6])

否则，由于张量仅包含正值，您可以对列和掩码求和：

>>> x.sum(1).nonzero()[:,0]
tensor([1, 3, 5, 6])

在 pytorch 张量中，return 特定值行的索引数组

In a pytorch tensor, return an array of indices of the rows of specific value

python

indexing

pytorch