python - Torch.sort 和 argsort 在相同元素的情况下随机排序

Question

当遇到相同的元素时，torch.sort 和 argsort 以随机方式对张量进行排序。在 numpy 中不是这种情况。我有一个已经根据第二列排序的元素列表，现在我想使用第一列对它进行排序，但保留较早的排序以防新排序中出现并列。

import torch

a = torch.tensor(
        [[ 0., 3.],
        [ 2., 3.],
        [ 2., 2.],
        [10., 2.],
        [ 0., 2.],
        [ 6., 2.],
        [10., 1.],
        [ 2., 1.],
        [ 0., 1.],
        [ 6., 1.],
        [10., 0.],
        [12., 0.]]
)
print(a[torch.argsort(a[:, 0])])

输出:

tensor([[ 0.,  3.],
        [ 0.,  2.],
        [ 0.,  1.],
        [ 2.,  1.],
        [ 2.,  2.],
        [ 2.,  3.],
        [ 6.,  1.],
        [ 6.,  2.],
        [10.,  1.],
        [10.,  2.],
        [10.,  0.],
        [12.,  0.]])

NumPy 的:

import numpy as np

a = np.array(
        [[ 0., 3.],
        [ 2., 3.],
        [ 2., 2.],
        [10., 2.],
        [ 0., 2.],
        [ 6., 2.],
        [10., 1.],
        [ 2., 1.],
        [ 0., 1.],
        [ 6., 1.],
        [10., 0.],
        [12., 0.]]
)
print(a[np.argsort(a[:, 0])])

输出:

[[ 0.  3.]
 [ 0.  2.]
 [ 0.  1.]
 [ 2.  3.]
 [ 2.  2.]
 [ 2.  1.]
 [ 6.  2.]
 [ 6.  1.]
 [10.  2.]
 [10.  1.]
 [10.  0.]
 [12.  0.]]

这可能是什么原因？我能做些什么来避免它？

Answer 1

根据 torch 1.9.0，您可以使用选项 stable=True 运行排序。参见 https://pytorch.org/docs/1.9.0/generated/torch.sort.html?highlight=sort#torch.sort

>>> x = torch.tensor([0, 1] * 9)
>>> x.sort()
torch.return_types.sort(
    values=tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1]),
    indices=tensor([ 2, 16,  4,  6, 14,  8,  0, 10, 12,  9, 17, 15, 13, 11,  7,  5,  3,  1]))
>>> x.sort(stable=True)
torch.return_types.sort(
    values=tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1]),
    indices=tensor([ 0,  2,  4,  6,  8, 10, 12, 14, 16,  1,  3,  5,  7,  9, 11, 13, 15, 17]))

文档说这只在 CPU 上，但它很快就会出现在 GPU 排序上，因为该文档警告已在 github 的 master 分支中删除(根据 https://github.com/pytorch/pytorch/pull/61685)