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函数找到最大的两个值及其索引和最小的两个值及其索引,并打印出结果。

相关推荐
灏瀚星空10 分钟前
从基础到实战的量化交易全流程学习:1.2 金融市场基础
笔记·python·信息可视化·系统架构·开源
Se7en25814 分钟前
AI 乱写代码怎么破?使用 Context7 MCP Server 让 AI 写出靠谱代码!
人工智能
零零刷18 分钟前
德州仪器(TI)—TDA4VM芯片详解(1)—产品特性
人工智能·嵌入式硬件·深度学习·神经网络·自动驾驶·硬件架构·硬件工程
用户277844910499326 分钟前
Python打造Excel记账模板,摸鱼时间也能轻松理财
人工智能·python
阿杜杜不是阿木木26 分钟前
03.使用spring-ai玩转MCP
java·人工智能·spring boot·spring·mcp·spring-ai
IT古董27 分钟前
【漫话机器学习系列】230.模型复杂度对训练和测试误差的影响(The Effect Of Model Complexity Training And Test)
人工智能
奕川28 分钟前
深入浅出:MCP 协议及应用
人工智能·后端
IT古董29 分钟前
【漫话机器学习系列】233.激活阈(Threshold Activation)
人工智能
极客智谷31 分钟前
MCP协议与智能体开发:构建动态协作AI系统的基石
人工智能
极客智谷31 分钟前
Spring AI应用系列——基于DeepSeek客户端的Chat应用
人工智能·后端