生成一维张量作为二维张量行的唯一索引答案

【问题标题】：generate 1D tensor as unique index of rows of an 2D tensor生成一维张量作为二维张量行的唯一索引
【发布时间】：2022-06-20 23:02:41
【问题描述】：

假设我们通过为每个不同的行赋予不同的索引（从 0 到 the number of rows - 1）将 2D 张量转换为 1D 张量。

[[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 上实现？

【问题讨论】：

标签： python pytorch tensor

【解决方案1】：

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

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

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

【讨论】：

【解决方案2】：

您可以重新调整要检查的张量，然后在每一行中检查所有值都是 True 并返回带有 argwhere 的索引，如下所示：

tns0 = torch.tensor([[1,2],[1,3],[1,4]])
tns1 = torch.tensor([[1,2],[1,2],[1,4]])
tns2 = torch.tensor([[1,2],[1,3],[1,2]])

a = torch.all(torch.reshape(tns1, (-1,1,2)) == tns0, dim=2)
torch.argwhere(a)[:,1]
# tensor([0, 0, 2])

b = torch.all(torch.reshape(tns2, (-1,1,2)) == tns0, dim=2)
torch.argwhere(b)[:,1]
#tensor([0, 1, 0])

说明：（如果你在每一行找到True，你就可以找到你想要的索引）

>>> c = torch.reshape(tns1, (-1,1,2)) == tns0
>>> c
tensor([[[ True,  True],
         [ True, False],
         [ True, False]],

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

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


>>> torch.all(c, dim=2)
tensor([[ True, False, False],
        [ True, False, False],
        [False, False,  True]])

【讨论】：

【解决方案3】：

Ivan 的想法奏效了。发现使用torch.unique_consecutive更好地保持行的原始顺序。

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

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

【讨论】：