python - Pytorch 张量 - 如何通过特定张量获取索引

Question

我有张量

t = torch.tensor([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]])

和一个查询张量

q = torch.tensor([1, 0, 0, 0])

有没有办法获得 q 的索引？喜欢

indexes = t.index(q) # get back [0, 3]

在 pytorch 中？

Answer 1

怎么样

In [1]: torch.nonzero((t == q).sum(dim=1) == t.size(1))
Out[1]: 
tensor([[ 0],
        [ 3]])

比较 t == q在 t 之间执行逐元素比较和 q ，由于您正在寻找整行匹配，您需要 .sum(dim=1)沿着行，看看哪一行是完美匹配的 == t.size(1) .

从 v0.4.1 开始， torch.all() 支持 dim争论:

torch.all(t==q, dim=1)