torch.argsort 函数介绍

torch.argsort 是 PyTorch 中用于返回张量沿指定维度的排序索引的函数。它不会直接对张量的值进行排序,而是返回一个与张量相同形状的索引张量,指示出原张量中每个元素的排序顺序。

函数签名

复制代码
torch.argsort(input, dim=-1, descending=False)

参数

  • input:输入的张量。
  • dim :指定进行排序的维度,默认为最后一个维度(-1)。
  • descending :如果为 True,则返回按降序排列的索引;如果为 False(默认),则按升序排列。

返回

返回与输入张量 input 形状相同的张量,其中每个元素为排序后的索引值。

使用场景

  • 需要知道张量中元素的排序顺序而不是排序后的实际值时,可以使用 argsort
  • 可以用于生成一种排序的掩码,比如对某些值进行有序操作。

示例代码

1. 基本用法
复制代码
import torch

# 创建一个张量
x = torch.tensor([3.5, 1.2, 4.8, 2.9])

# 对张量进行升序排序,并返回排序后的索引
sorted_indices = torch.argsort(x)
print(sorted_indices)  # 输出:tensor([1, 3, 0, 2])

解释:torch.argsort 返回的是按升序排列元素的索引。索引 1 对应值 1.2(最小),依次类推。

2. 在多维张量上使用
复制代码
import torch

# 创建一个 2D 张量
x = torch.tensor([[4, 1, 3],
                  [2, 8, 5]])

# 在每一行上对张量进行升序排序,dim=1 表示对行操作
sorted_indices = torch.argsort(x, dim=1)
print(sorted_indices)

# 在每一列上对张量进行升序排序,dim=0 表示对列操作
sorted_indices_col = torch.argsort(x, dim=0)
print(sorted_indices_col)

输出

复制代码
tensor([[1, 2, 0],
        [0, 2, 1]])

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

解释:第一个例子返回的是按行升序排列的索引,第二个例子返回的是按列升序排列的索引。

3. 降序排序
复制代码
import torch

# 创建一个张量
x = torch.tensor([3.5, 1.2, 4.8, 2.9])

# 对张量进行降序排序,并返回排序后的索引
sorted_indices = torch.argsort(x, descending=True)
print(sorted_indices)  # 输出:tensor([2, 0, 3, 1])

解释:这里使用了 descending=True,所以 torch.argsort 按降序返回索引,索引 2 对应值 4.8(最大),依次类推。

4. 与 gather 配合使用

你可以使用 torch.argsort 来获取排序后的索引,并结合 torch.gather 获取排序后的值:

复制代码
import torch

# 创建一个张量
x = torch.tensor([3.5, 1.2, 4.8, 2.9])

# 获取排序后的索引
sorted_indices = torch.argsort(x)

# 根据索引获取排序后的值
sorted_x = torch.gather(x, 0, sorted_indices)
print(sorted_x)  # 输出:tensor([1.2000, 2.9000, 3.5000, 4.8000])

解释:通过 torch.gather 函数,可以根据排序后的索引来重新排列原张量的值。

总结

torch.argsort 函数非常适合需要知道张量元素排序顺序的场景。它返回的是元素在升序或降序排序中的索引,而不是排序后的实际值。结合其他 PyTorch 函数(如 gather),可以灵活实现很多数据处理操作。

相关推荐
胖哥真不错23 分钟前
Python实现NOA星雀优化算法优化随机森林回归模型项目实战
python·机器学习·项目实战·随机森林回归模型·noa星雀优化算法
编程咕咕gu-24 分钟前
从零开始玩python--python版植物大战僵尸来袭
开发语言·python·python基础·pygame·python教程
代码的乐趣1 小时前
支持selenium的chrome driver更新到135.0.7049.42
chrome·python·selenium
带你去吃小豆花1 小时前
在亚马逊云科技上使用n8n快速构建个人AI NEWS助理
人工智能·科技·ai·云原生·aws
AI技术控2 小时前
计算机视觉算法实战——基于YOLOv8的自动驾驶障碍物实时感知系统
人工智能·yolo·自动驾驶
SsummerC4 小时前
【leetcode100】数组中的第K个最大元素
python·算法·leetcode
伊玛目的门徒4 小时前
解决backtrader框架下日志ValueError: I/O operation on closed file.报错(jupyternotebook)
python·backtrader·量化·日志管理·回测
java1234_小锋5 小时前
一周学会Pandas2 Python数据处理与分析-编写Pandas2 HelloWord项目
python·pandas·python数据分析·pandas2
凯强同学6 小时前
第十四届蓝桥杯大赛软件赛省赛Python 大学 C 组:7.翻转
python·算法·蓝桥杯
耘瞳科技6 小时前
喜讯 | 耘瞳科技视觉检测与测量装备荣膺“2024机器视觉创新产品TOP10”
人工智能·科技·视觉检测