python - 删除 torch.tensor 中的重复行

Question

我有一个形状为 (n,m) 的 torch.tensor，我想删除重复的行(或至少找到它们)。例如:

t1 = torch.tensor([[1, 2, 3], [4, 5, 6], [1, 2, 3], [4, 5, 6]])
t2 = remove_duplicates(t1)

t2 现在应该等于 tensor([[1, 2, 3], [4, 5, 6]])，即行 1 和 3 被删除。你知道执行此操作的方法吗？

我想用 torch.unique 做点什么，但我不知道该怎么做。

Answer 1

您可以简单地利用 torch.unique 的参数 dim。

t1 = torch.tensor([[1, 2, 3], [4, 5, 6], [1, 2, 3], [4, 5, 6], [7, 8, 9]])
torch.unique(t1, dim=0)

这样你就可以得到你想要的结果:

tensor([[1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]])

Here您可以阅读该参数的含义。