在 pytorch 张量中，返回特定值行的索引数组答案

【问题标题】：In a pytorch tensor, return an array of indices of the rows of specific value在 pytorch 张量中，返回特定值行的索引数组
【发布时间】：2021-11-08 08:41:27
【问题描述】：

给定以下具有全零向量和具有一和零向量的张量：

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

我怎样才能有一个向量的索引数组，其中包含一个和零，所以输出是这样的：

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

更新

一种方法是：

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

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

【问题讨论】：

标签： python indexing pytorch

【解决方案1】：

另一种方法是使用torch.Tensor.any 加上torch.Tensor.nonzero：

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

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

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

【讨论】：