寻找两个pytorch张量的不交集

Question

提前感谢大家的帮助！我在 PyTorch 中尝试做的事情类似于 numpy 的setdiff1d. 例如给出以下两个张量：

t1 = torch.tensor([1, 9, 12, 5, 24]).to('cuda:0')
t2 = torch.tensor([1, 24]).to('cuda:0')

预期输出应为（已排序或未排序）：

torch.tensor([9, 12, 5])

理想情况下，操作是在 GPU 上完成的，GPU 和 CPU 之间没有来回。非常感激！

Answer 1

我遇到了同样的问题，但是在使用更大的阵列时，建议的解决方案太慢了。以下简单的解决方案适用于 CPU 和 GPU，并且比其他建议的解决方案快得多：

combined = torch.cat((t1, t2))
uniques, counts = combined.unique(return_counts=True)
difference = uniques[counts == 1]
intersection = uniques[counts > 1]

Answer 2

如果您不想离开 cuda，解决方法可能是：

t1 = torch.tensor([1, 9, 12, 5, 24], device = 'cuda')
t2 = torch.tensor([1, 24], device = 'cuda')
indices = torch.ones_like(t1, dtype = torch.uint8, device = 'cuda')
for elem in t2:
    indices = indices & (t1 != elem)  
intersection = t1[indices]