两个n维pytorch张量的非交集

Non-intersection of two n-dimentional pytorch tensors

提前感谢大家的帮助!我在 PyTorch 中尝试做的是为许多维度的张量计算非交集(我们称之为 torch.nonintersection)(没有 for 循环,因为我希望它在 GPU 上有效执行)。所以这是它应该如何工作的例子:

a = torch.tensor([[ 0.,  0.], [ 0.,  1.], [ 0.,  2.], [ 1.,  0.], [ 1.,  1.], [ 1.,  2.], [ 1.,  3.], 
                  [ 2.,  0.], [ 2.,  1.], [ 2.,  2.]])
b = torch.tensor([[ 2.,  0.], [ 2.,  1.], [ 2.,  2.], [ 1.,  0.], [ 1.,  1.], [ 1.,  2.], [ 1.,  3.]])

torch.spec_unique(a,b) = torch.tensor([ 0.,  0.], [ 0.,  1.], [ 0.,  2.])

我有 for 循环的类似物,但它们现在需要太多时间。任何想法如何完成?非常感谢!

ab 的形状在您的情况下不同(a(10, 2),而 b(7, 2)),所以你必须将较大的张量“切割”成较小的张量。下面的函数用简单的 if 处理它(它不会减慢你的计算速度):

def non_intersection(a, b):
    if a.shape[0] > b.shape[0]:
        return torch.nonzero(a[: b.shape[0]] != b, as_tuple=False)
    return torch.nonzero(b[: a.shape[0]] != a, as_tuple=False)

这会 return:

tensor([[0, 0],
        [1, 0],
        [2, 0]])

所以列的顺序相反。如果您希望像您的示例中那样获得列,您可以在 non_intersection 输出上执行此操作:

torch.index_select(non_intersection(a, b), 1, torch.tensor([1, 0])))