torch.argsort

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

沿指定维度,按值的升序返回张量的索引。

这是 torch.sort() 返回的第二个值。请参阅该函数的文档以了解其确切含义。

如果 stableTrue,排序过程会保持相等元素的原有顺序。若为 False,相等值之间的相对顺序无法保证。True 模式较慢。

参数
  • input (Tensor) – 需要输入的张量。

  • dim (int, 可选) – 需要进行排序的维度

  • descending (bool, 可选) – 控制是否按升序或降序排列

  • stable (bool, 可选) – 控制相等元素的相对顺序

示例:

>>> a = torch.randn(4, 4)
>>> a
tensor([[ 0.0785,  1.5267, -0.8521,  0.4065],
        [ 0.1598,  0.0788, -0.0745, -1.2700],
        [ 1.2208,  1.0722, -0.7064,  1.2564],
        [ 0.0669, -0.2318, -0.8229, -0.9280]])


>>> torch.argsort(a, dim=1)
tensor([[2, 0, 3, 1],
        [3, 2, 1, 0],
        [2, 1, 0, 3],
        [3, 2, 1, 0]])
本页目录