Python PyTorch argsort用法及代码示例

本文简要介绍python语言中 torch.argsort 的用法。

用法:

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

参数：

• input(Tensor) -输入张量。

• dim(int,可选的) -要排序的维度

• descending(bool,可选的) -控制排序顺序(升序或降序)

返回沿给定维度按值升序对张量进行排序的索引。

这是 torch.sort() 返回的第二个值。有关此方法的确切语义，请参阅其文档。

例子：

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