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]])