torch.sort()和torch.argsort()简要介绍

定义

• torch.sort(input,dim,descending)
• torch.argsort(input,dim,descending)

用法

• torch.sort：对输入数据排序，返回两个值，即排序后的数据values和其在原矩阵中的坐标indices
• torch.argsort：同torch.sort()返回的indices

参数

• input：输入矩阵
• dim：排序维度，默认为dim=1,即对行排序
• descending：排序方式（从小到大和从大到小），默认为从小到大排序（即descending=False)

示例

torch.sort()

import torch

a = torch.tensor([[2,3,1],[0,5,6]])

print(torch.sort(a))
print(torch.sort(a,dim=0))
print(torch.sort(a,dim=0,descending=True))

>>>torch.return_types.sort(
values=tensor([[1, 2, 3],
               [0, 5, 6]]),
indices=tensor([[2, 0, 1],
                [0, 1, 2]]))

torch.return_types.sort(
values=tensor([[0, 3, 1],
               [2, 5, 6]]),
indices=tensor([[1, 0, 0],
                [0, 1, 1]]))

torch.return_types.sort(
values=tensor([[2, 5, 6],
               [0, 3, 1]]),
indices=tensor([[0, 1, 1],
                [1, 0, 0]]))
​

torch.argsort()

print(torch.argsort(a))
print(torch.argsort(a,dim=0 ))
print(torch.argsort(a,dim=0,descending = False))

>>>tensor([[2, 0, 1],
           [0, 1, 2]])
tensor([[1, 0, 0],
        [0, 1, 1]])
tensor([[0, 1, 1],
        [1, 0, 0]])