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

Question

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

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

Answer 1

这是一种方法：

使用 numpy.argsort()
使用torch.from_numpy()
将结果转换为张量
import torch import numpy as np A = [0,1,2,3,0,0,1,1,2,2,3,3] x = np.array(A) y = torch.from_numpy(np.argsort(x, kind='mergesort')) print(y)

Answer 2

这是一个利用 broadcasting, torch.unique(), and torch.nonzero() 的纯基于 PyTorch 的解决方案。这将是一个巨大的提升，特别是对于基于 implementation/run 的 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-based解决方案不同，这里我们不用担心我们kind的排序算法' ve 使用，因为我们根本没有使用任何排序。

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

Pytorch argsort ordered, with duplicate elements in the tensor

python

sorting

machine-learning

pytorch

tensor