时间戳

首页 > 解决方案 > Pytorch argsort 有序，张量中有重复元素

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

问题描述

我有一个向量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 循环的情况下解决这个问题？这样的操作是我的 NN 模型的一部分，如果没有有效地完成会导致性能问题。谢谢！

标签： pythonsortingmachine-learningpytorchtensor

解决方案

这是一个纯粹的基于 PyTorch 的解决方案，利用broadcasting、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)

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我们必须使用什么排序算法，因为我们根本没有使用任何排序。

推荐阅读