python - 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

这是一个利用 broadcasting 的纯 PyTorch 解决方案, torch.unique() , 和 torch.nonzero() .这将有很大的插入作用，特别是对于基于 GPU 的实现/运行而言，如果我们必须切换回 NumPy，这是不可能的，argsort 然后转移回 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)

一旦我们有了这个 bool 张量，我们就可以通过在转置 bool 张量后检查有 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 排序算法，因为我们根本没有使用任何排序。