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

【问题标题】：Pytorch tensor - How to get the indexes by a specific tensorPytorch 张量 - 如何通过特定张量获取索引
【发布时间】：2019-01-13 04:45:11
【问题描述】：

我有一个张量

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

有没有办法获取qlike的索引

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

在 pytorch 中？

【问题讨论】：

标签： python pytorch

【解决方案1】：

请试试这个，我没有在这台电脑上安装手电筒。

import torch

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

index = torch.nonzero(torch.sum((t == q), dim=1) == t.shape[1])

编辑说明：针对 Shai 提出的问题进行了编辑。

【讨论】：

您的代码执行 element-wise 比较，得到[[ 0, 0], [ 0, 1], [ 0, 2], [ 0, 3], [ 1, 1], [ 1, 3], [ 2, 2], [ 2, 3], [ 3, 0], [ 3, 1], [ 3, 2], [ 3, 3]]。您需要改为比较 rows。
@Shai Ic，可以轻松修复

【解决方案2】：

怎么样

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

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

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

【讨论】：

【解决方案3】：

更多的 pytorch 原生方法是：

torch.all(q.repeat((t.shape[1],1))==t, dim=1)

【讨论】：