生成 1D 张量作为 2D 张量的行的唯一索引

假设我们通过为每个不同的行提供不同的索引来将 2D 张量转换为 1D 张量，从0 to the number of rows - 1.

[[1,2],[1,3],[1,4]] -> [0,1,2]

但如果有相同的行，那么我们重复索引，如下所示。

[[1,2],[1,2],[1,4]] -> [0,0,2]

[[1,2],[1,3],[1,2]] -> [0,1,0]

如何在 PyTorch 上实现这个？

您可以使用torch.Tensor.unique并返回提供开箱即用索引的逆：

>>> _, i = x.unique(dim=0, return_inverse=True)

>>> i # first example
tensor([0, 1, 2])