Pytorch argsort 有序，张量中有重复元素

Question

我有一个向量A = [0,1,2,3,0,0,1,1,2,2,3,3]。我需要以递增的方式对它进行排序，以便它以有序的方式列出，并从中提取 argsort。为了更好地解释这一点，我需要对 A 进行排序，使其返回B = [0,4,5,1,6,7,2,8,9,3,10,11]。但是，当我使用 pyotrch 时，torch.argsort(A)它会返回B = [4,5,0,1,6,7,2,8,9,3,10,11].

我假设这样做的算法无法由我来控制。有没有办法在不引入 for 循环的情况下解决这个问题？此类操作是我的神经网络模型的一部分，如果不能有效完成，将会导致性能问题。谢谢！

Answer 1

这是一个纯粹基于 PyTorch 的解决方案broadcasting，利用torch.unique()、 和torch.nonzero()。argsort这将带来巨大的提升，特别是对于基于 GPU 的实现/运行来说，如果我们必须切换回 NumPy，然后再转移回 PyTorch（如其他方法中所建议的那样），这是不可能的。

# our input tensor
In [50]: A = torch.tensor([0,1,2,3,0,0,1,1,2,2,3,3])

# construct an intermediate boolean tensor
In [51]: boolean = A[:, None] == torch.unique(A)

In [52]: boolean
Out[52]: 
tensor([[1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 1],
        [1, 0, 0, 0],
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 1],
        [0, 0, 0, 1]], dtype=torch.uint8)

1一旦我们有了这个布尔张量，我们就可以通过检查转置布尔张量后存在的位置来找到所需的索引。

这将为我们提供排序input和indices. 由于我们只想要索引，因此我们可以通过为最后一列（1或-1）建立索引来获取它们

In [53]: torch.nonzero(boolean.t())[:, -1]
Out[53]: tensor([ 0,  4,  5,  1,  6,  7,  2,  8,  9,  3, 10, 11])

这是 OP 在评论中提供的另一个示例的结果：

In [55]: A_large = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9])

In [56]: boolean_large = A_large[:, None] == torch.unique(A_large)

In [57]: torch.nonzero(boolean_large.t())[:, -1]
Out[57]: 
tensor([ 0, 10, 11,  1, 12, 13,  2, 14, 15,  3, 16, 17,  4, 18, 19,  5, 20, 21,
         6, 22, 23,  7, 24, 25,  8, 26, 27,  9, 28, 29])

注意：与其他答案中提出的基于 NumPy 的解决方案不同，这里我们不必担心kind我们必须使用什么排序算法，因为我们根本不使用任何排序。