获取张量 a 中存在于张量 b 中的元素的索引答案

【问题标题】：Get indices of elements in tensor a that are present in tensor b获取张量 a 中存在于张量 b 中的元素的索引
【发布时间】：2020-03-29 17:41:49
【问题描述】：

例如，我想获取张量 a 中值 0 和 2 的元素的索引。这些值（0 和 2）存储在张量 b 中。我已经设计了一种 pythonic 方式来做到这一点（如下所示），但我不认为列表推导被优化为在 GPU 上运行，或者也许有一种我不知道的更 PyTorchy 方式来做到这一点。

import torch
a = torch.tensor([0, 1, 0, 1, 1, 0, 2])
b = torch.tensor([0, 2])
torch.tensor([x in b for x in a]).nonzero()

>>>> tensor([[0],
             [2],
             [5],
             [6]])

还有其他建议吗？或者这是可以接受的方式吗？

【问题讨论】：

在 numpy 中，我相信你可以使用 np.nonzero(np.isin(a, b)) 来做到这一点，假设 a 和 b 是 numpy 数组。不幸的是，isin 目前不是 pytorch 操作，尽管有一个开放的功能请求以及依赖于in the pytorch issue tracker 提供的广播的次优版本。

标签： python pytorch tensor

【解决方案1】：

这是一种更有效的方法（正如 jodag 在 cmets 中发布的链接中所建议的那样...）：

(a[..., None] == b).any(-1).nonzero()

【讨论】：

这很好。 a[..., None] 解压和转置向量a。 a[..., None] == b 将转置后的a 的每个元素与b 的所有元素进行比较，为每次比较生成一个布尔矩阵。 .any(-1) 使用适当的维度将比较汇总为一个向量，nonzero() 将布尔值转换为索引。
不应该像(a[..., None] == b).any(-1).nonzero().squeeze()吗？因为，否则，它会增加额外的维度。
@DmytroSytro 我认为答案给出的解决方案与我在原始问题中发布的完全一样，但是是的，在这两种情况下都会创建一个冗余维度，squeeze() 摆脱了它。