如何在 pytorch 中按第一维度对张量进行排序？

Question

我有一个二维张量，我想像这个例子一样按第一维排序。

a = torch.FloatTensor(
   [[5, 5],
    [5, 3],
    [3, 5],
    [6, 4],
    [3, 7]])

我希望排序后有这样的结果

a = torch.FloatTensor(
   [[3, 5],
    [3, 7],
    [5, 3],
    [5, 5],
    [6, 4]])

在pytorch中可以这样做吗? 我知道在numpy中是可以做到的，但是我想在GPU中用Torch来做。

Answer 1

按第一列排序，然后使用指数对整个数组进行排序。

a[a[:, 0].sort()[1]]

输出：

tensor([[3., 5.],
        [3., 7.],
        [5., 5.],
        [5., 3.],
        [6., 4.]])

如果你真的需要交错排列的话:

b = a[a[:, 1].sort()[1]]
b[b[:, 0].sort()[1]]

输出：

tensor([[3., 5.],
        [3., 7.],
        [5., 3.],
        [5., 5.],
        [6., 4.]])

Answer 2

torch.stack(sorted(a, key=lambda a: a[0]))

输出结果是：

tensor([[3., 5.],
        [3., 7.],
        [5., 5.],
        [5., 3.],
        [6., 4.]])