如何获取两个不同大小的 PyTorch 张量中相等元素的索引？

【问题标题】：How to get the indexes of equal elements in two different size PyTorch tensors?如何获取两个不同大小的 PyTorch 张量中相等元素的索引？
【发布时间】：2021-02-28 13:09:06
【问题描述】：

假设我有两个 PyTorch 张量：

t_1d = torch.Tensor([6, 5, 1, 7, 8, 4, 7, 1, 0, 4, 11, 7, 4, 7, 4, 1])
t = torch.Tensor([4, 7])

我想获取集合之间完全匹配交集的索引张量 t_1d 与张量 t。

t_1d 和 t 的期望输出：[5, 12]（精确交集的第一个索引）

最好在 GPU 上用于大型张量，因此没有循环或 Numpy 强制转换。

【问题讨论】：

标签： python pytorch

【解决方案1】：

一般来说，我们可以检查t 中的每个元素与t_1d 中的元素相等的位置。

之后，将最后一个元素从第一个元素移回原来的位置（一般情况下，这里移-1）并检查数组是否相等：

intersection = (t_1d == t[0]) & torch.roll(t_1d == t[1], shifts=-1)
torch.where(intersection)[0] # torch.tensor([5, 12])

【讨论】：