pytorch torch.topk函数介绍

在 PyTorch 中,torch.topk函数用于在输入张量中找到最大的k个值及其索引。

一、函数语法

torch.topk(input, k, dim=None, largest=True, sorted=True, out=None)

  • input:输入张量。
  • k:要返回的最大或最小元素的数量。
  • dim(可选):要进行操作的维度。如果为None,则在扁平的输入张量上进行操作。
  • largest(可选):如果为True,则返回最大的k个值;如果为False,则返回最小的k个值。
  • sorted(可选):如果为True,则返回的k个值将按降序(如果largest=True)或升序(如果largest=False)排列;如果为False,则返回的k个值的顺序是未定义的。
  • out(可选):输出张量,可以是一个已存在的张量,用于存储结果。

二、返回值

该函数返回一个包含两个张量的元组:

  1. 第一个张量是包含最大或最小的k个值的张量。
  2. 第二个张量是包含这些值在输入张量中的索引的张量。

三、使用示例

复制代码
import torch

# 创建一个二维张量
tensor = torch.tensor([[4, 2, 3], [1, 5, 6]])

# 找到每行中的最大的两个值及其索引
values, indices = torch.topk(tensor, k=2, dim=1, largest=True)
print("最大的两个值:", values)
print("对应的索引:", indices)

# 找到每列中的最小的两个值及其索引
values, indices = torch.topk(tensor, k=2, dim=0, largest=False)
print("最小的两个值:", values)
print("对应的索引:", indices)

在上述示例中,首先创建了一个二维张量。然后,分别在行维度和列维度上使用torch.topk函数找到最大的两个值及其索引和最小的两个值及其索引,并打印出结果。

相关推荐
递归不收敛20 小时前
吴恩达机器学习课程(PyTorch适配)学习笔记:1.5 决策树与集成学习
pytorch·学习·机器学习
菜鸟‍20 小时前
【论文学习】2025年图像处理顶会论文
图像处理·人工智能·学习
飞哥数智坊20 小时前
“成章”开源首更:3大体验优化,合集管理补齐!
人工智能·开源·ai编程
GoldenSpider.AI20 小时前
划时代的技术飞跃:OpenAI DevDay 2025 全面深度解读
人工智能·chatgpt·codex·openai devday·gpt-5·sam altman·sora 2
Qiuner20 小时前
【机器学习】(一)实用入门指南——如何快速搭建自己的模型
人工智能·机器学习
whaosoft-14320 小时前
51c视觉~3D~合集7
人工智能
AI视觉网奇20 小时前
开源3d数字人学习笔记2025
人工智能·计算机视觉
IT_陈寒21 小时前
Java性能优化:3个90%开发者都忽略的高效技巧,让你的应用提速50%!
前端·人工智能·后端
無斜21 小时前
【LabVIEW实用开发】--- LabVIEW调用python脚本
开发语言·python·labview
北京阿法龙科技有限公司21 小时前
AI 驱动的 AR眼镜巡检技术方案:让工业缺陷识别更精准高效|阿法龙XR云平台
人工智能·ar·xr