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

Pytorch tensor - How to get the indexes by a specific tensor(Pytorch 张量 - 如何通过特定张量获取索引)

问题描述

我有一个张量

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

有没有办法像

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

在pytorch中?

推荐答案

怎么样

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

Comparing t == q 在 t 和 q 之间执行 element-wise 比较，因为您正在寻找对于整行匹配，您需要沿行 .sum(dim=1) 并查看哪一行是完美匹配 == t.size(1).

Comparing t == q performs element-wise comparison between t and q, since you are looking for entire row match, you need to .sum(dim=1) along the rows and see what row is a perfect match == t.size(1).

从 v0.4.1 开始，torch.all() 支持 dim 参数:

As of v0.4.1, torch.all() supports dim argument:

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

这篇关于Pytorch 张量 - 如何通过特定张量获取索引的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持编程学习网！