我是Pytorch的新手。给定一个张量集，我需要根据键值对这些张量进行排序。例如，

A = [[0.9133, 0.5071, 0.6222, 3.], [0.5951, 0.9315, 0.6548, 1.], [0.7704, 0.0720, 0.0330, 2.]]

排序后的预期结果是：

A' = [[0.5951, 0.9315, 0.6548, 1.], [0.7704, 0.0720, 0.0330, 2.], [0.9133, 0.5071, 0.6222, 3.]]

我尝试使用Python中的sorted函数，但这在我的训练过程中非常耗时。如何更有效地实现这一点？谢谢！

回答：

%%timeit -r 10 -n 10A[A[:,-1].argsort()]38.6 µs ± 23 µs per loop (mean ± std. dev. of 10 runs, 10 loops each)%%timeit -r 10 -n 10sorted(A, key = lambda x: x[-1])69.6 µs ± 34.8 µs per loop (mean ± std. dev. of 10 runs, 10 loops each)

两者输出

tensor([[0.5951, 0.9315, 0.6548, 1.0000],        [0.7704, 0.0720, 0.0330, 2.0000],        [0.9133, 0.5071, 0.6222, 3.0000]])

然后是

%%timeit -r 10 -n 10a, b = torch.sort(A, dim=-2)最慢的一次运行比最快的一次运行慢了8.45倍。这可能意味着中间结果被缓存了。14.3 µs ± 18.1 µs per loop (mean ± std. dev. of 10 runs, 10 loops each)

其中a是排序后的张量，b是索引