torch.argsort用法及利用torch.argsort恢复张量的原始顺序

torch.argsort

PyTorch函数，用于对张量中的元素进行排序并返回排序后的索引。它不直接返回排序后的元素，而是返回这些元素在原始张量中的位置索引。

基本用法

torch.argsort的基本语法如下：

torch.argsort(input, dim=-1, descending=False)

• input：待排序的输入张量
• dim：要排序的维度。默认是最后一个维度
• descending：排序顺序，如果为True，则按降序排序；如果为False（默认值），则按升序排序

利用torch.argsort恢复张量的原始顺序

1、假设有一个为noise的张量，形状为[1, 3]，

noise = torch.tensor([[0.8, 0.3, 0.5]])

2、对noise执行torch.argsort(noise, dim=1)时，得到的ids_shuffle是每行元素按值排序后的索引,

ids_shuffle = torch.argsort(noise, dim=1)

"""
输出:
[[1, 2, 0]]
"""

这表示在noise的第一行中，最小的值是0.3（位于原始索引1），其次是0.5（位于原始索引2），最大的是0.8（位于原始索引0）

3、随后，对ids_shuffle应用torch.argsort，实际上是在对每行的索引进行排序。这次排序是基于索引的值进行的。ids_restore的结果告诉我们要如何通过ids_shuffle提供的索引来恢复原始数组的顺序，

ids_restore = torch.argsort(ids_shuffle, dim=1)

"""
输出:
[[2, 0, 1]]
"""

这个结果告诉我们，

• ids_restore中的第一个值为2，对应的是ids_shuffle第三个位置的元素（0），而0则表示原始张量中的0.8
• ids_restore中的第二个值为0，对应的是ids_shuffle第一个位置的元素（1），而1则表示原始张量中的0.3
• ids_restore中的第三个值为1，对应的是ids_shuffle第二个位置的元素（2），而2则表示原始张量中的0.5

4、利用torch.gather恢复张量的原始顺序，

shuffled_data = torch.gather(noise, 1, ids_shuffle)
"""
ids_shuffle = tensor([[1, 2, 0]])
shuffled_data = tensor([[0.3000, 0.5000, 0.8000]])
"""

restored_data = torch.gather(shuffled_data, 1, ids_restore)
"""
ids_restore = tensor([[2, 0, 1]])
restored_data = tensor([[0.8000, 0.3000, 0.5000]])
"""