有没有比 for 循环更快的方法来比较两个张量中的两行？

Question

我是 pytorch 的新手，我想检查 tensor1 中的每个元素是否属于 tensor2 中同一行的一部分，例如，如果我有两个张量如下：

a = torch.tensor([[0,5],[1,5], [4,5], [7,8]])
b = torch.tensor([[0,1],[2,3], [7,5], [-1,7]])

输出应如下所示：

[ True, False]
[False, False]
[False, False]
[False, False]
[False, False]
[False,  True]
[False,  True]
[False, False]

我知道我可以使用一个简单的 for 循环来完成，但是在 pytorch 中有没有更快的方法？

Answer 1

你需要对两个张量使用 repeat_interleave 然后 reshape 第一个张量，然后你将它与另一个张量进行比较像这样

res = a.repeat_interleave(2, dim=1).reshape(-1, 2) == b.repeat_interleave(2, dim=0)

输出应该如下

output
tensor([[ True, False],
        [False, False],
        [False, False],
        [False, False],
        [False, False],
        [False,  True],
        [False,  True],
        [False, False]])

输出是张量。如果你需要它而不用 [] 你可以使用 flatten 方法

有没有比 for 循环更快的方法来比较两个张量中的两行？

is there a faster way to compare two raws in a two tensors than for loop?

pytorch

tensor