torch.sort()用法

torch.sort()用法

torch.sort(input, dim=- 1, descending=False, stable=False, out=None)

这个函数用于将输入的tensor按照指定维度进行排序,并返回一个元组,元组里面的两个元素分别是(排序后的值, 对应的切片索引)

参数

• input(Tensor) 输入的张量
• dim(int, 可选择) 如果给定了就按照指定的维度进行排序,默认的是最后一个维度
• descending(bool , 可选择) 如果为TRUE就是降序排列,默认为False即升序排列
• stable(bool, 可选择) 如果为True那么如果张量中有两个元素相等,在排序之后返回的索引依旧会保持它们原先的相对位置,如果为False,则不一定会保持原先的相对位置
• out(tuple, 可选择)

维度操作的方向参见下图

示例

x=torch.randn(3,4)
x
>>>tensor([[-0.0454, -0.8172, -1.5806, -0.6127],
           [ 1.1915, -1.1318, -0.5910,  0.3758],
           [ 0.1943,  0.3682, -2.4066,  1.0772]])

sorted, indices = torch.sort(x) # 在维度1上按照升序排列
sorted
>>>tensor([[-1.5806, -0.8172, -0.6127, -0.0454],
           [-1.1318, -0.5910,  0.3758,  1.1915],
           [-2.4066,  0.1943,  0.3682,  1.0772]])
indices
>>>tensor([[2, 1, 3, 0],
          [1, 2, 3, 0],
          [2, 0, 1, 3]])

sorted, indices = torch.sort(x, 0) # 在维度0上按照升序排列
sorted
>>>tensor([[-0.0454, -1.1318, -2.4066, -0.6127],
           [ 0.1943, -0.8172, -1.5806,  0.3758],
           [ 1.1915,  0.3682, -0.5910,  1.0772]])
indices
>>>tensor([[0, 1, 2, 0],
           [2, 0, 0, 1],
           [1, 2, 1, 2]])

x = torch.tensor([0, 1] * 9)
x
>>>tensor([0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1])
# 不一定保持原先相对位置
x.sort()
>>>torch.return_types.sort(
    values=tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1]),
    indices=tensor([ 2, 16,  4,  6, 14,  8,  0, 10, 12,  9, 17, 15, 13, 11,  7,  5,  3,  1]))
# 一定保持原先相对位置
x.sort(stable=True)
>>>torch.return_types.sort(
    values=tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1]),
    indices=tensor([ 0,  2,  4,  6,  8, 10, 12, 14, 16,  1,  3,  5,  7,  9, 11, 13, 15, 17]))