YOLO26如何导出ONNX模型?推理格式转换详细步骤
在深度学习部署过程中,模型格式的兼容性至关重要。ONNX(Open Neural Network Exchange)作为一种开放的模型交换格式,能够实现跨框架、跨平台的模型部署,广泛应用于边缘设备、服务器推理引擎(如ONNX Runtime、TensorRT)等场景。
本文将基于最新YOLO26 官方版训练与推理镜像环境,详细介绍如何将训练好的 YOLO26 模型导出为 ONNX 格式,并完成推理验证,帮助开发者快速实现模型格式转换与工程化落地。
1. 镜像环境说明
本镜像基于YOLO26 官方代码库构建,预装了完整的深度学习开发环境,集成了训练、推理及评估所需的所有依赖,开箱即用。
- 核心框架:
pytorch == 1.10.0 - CUDA版本:
12.1 - Python版本:
3.9.5 - 主要依赖:
torchvision==0.11.0,torchaudio==0.10.0,cudatoolkit=11.3,numpy,opencv-python,pandas,matplotlib,tqdm,seaborn等。
该环境已预置ultralytics库及其相关工具链,支持直接调用model.export()方法进行模型导出,无需额外安装复杂依赖。
2. 导出ONNX模型的完整流程
### 2.1 准备工作:激活环境与进入项目目录
首先,启动镜像后需激活 Conda 环境并切换至项目目录:
conda activate yolo cd /root/workspace/ultralytics-8.4.2确保当前路径下存在已训练好的权重文件(如yolo26n.pt),或使用镜像内置的预下载权重。
### 2.2 使用Ultralytics API导出ONNX模型
YOLO26 提供了简洁的.export()接口,可一键完成 PyTorch 到 ONNX 的转换。
✅ 示例代码:export_onnx.py
创建一个新脚本export_onnx.py,内容如下:
# -*- coding: utf-8 -*- """ @File :export_onnx.py @Desc :将 YOLO26 模型导出为 ONNX 格式 """ from ultralytics import YOLO if __name__ == '__main__': # 加载模型(可以是 .pt 或 .yaml) model = YOLO(model='yolo26n.pt') # 支持 nano/small/medium/large 等变体 # 导出为 ONNX 格式 success = model.export( format='onnx', # 输出格式 dynamic=True, # 启用动态输入尺寸(推荐) simplify=True, # 简化计算图(优化推理性能) opset=12, # ONNX 算子集版本(建议 11~13) imgsz=640, # 输入图像大小(H=W) device='cuda' # 使用 GPU 加速导出(若可用) ) if success: print("✅ ONNX 模型导出成功!") else: print("❌ 模型导出失败,请检查日志。")🔍 参数详解:
| 参数 | 说明 |
|---|---|
format='onnx' | 指定输出格式为 ONNX |
dynamic=True | 允许输入尺寸动态变化(batch_size、height、width 可变) |
simplify=True | 调用onnx-simplifier工具简化模型结构,减少冗余节点 |
opset=12 | ONNX 算子集版本,兼容大多数推理引擎 |
imgsz=640 | 指定输入分辨率,默认正方形输入 |
⚠️ 注意:若未安装
onnx和onnxsim,请先执行:pip install onnx onnxsim
运行导出命令:
python export_onnx.py成功后将在当前目录生成yolo26n.onnx文件。
### 2.3 验证ONNX模型有效性
导出完成后,应验证 ONNX 模型是否结构正确、可加载、可推理。
✅ 示例代码:validate_onnx.py
# -*- coding: utf-8 -*- """ @File :validate_onnx.py @Desc :验证 ONNX 模型是否可正常加载和推理 """ import cv2 import numpy as np import onnxruntime as ort # 加载 ONNX 模型 session = ort.InferenceSession('yolo26n.onnx', providers=['CUDAExecutionProvider']) # 获取输入信息 input_name = session.get_inputs()[0].name input_shape = session.get_inputs()[0].shape print(f"Input Name: {input_name}") print(f"Input Shape: {input_shape}") # 读取测试图像 image = cv2.imread('./ultralytics/assets/zidane.jpg') ori_h, ori_w = image.shape[:2] # 图像预处理(resize + normalize + transpose) input_img = cv2.resize(image, (640, 640)) input_img = input_img.astype(np.float32) / 255.0 input_img = input_img.transpose(2, 0, 1) # HWC -> CHW input_img = np.expand_dims(input_img, axis=0) # NCHW # 执行推理 outputs = session.run(None, {input_name: input_img}) print(f"推理完成,输出数量: {len(outputs)}") # 查看第一个输出(通常是检测框) output = outputs[0] print(f"输出形状: {output.shape}") # 通常为 [1, num_boxes, 84](分类+坐标)🧪 输出示例:
Input Name: images Input Shape: [None, 3, 640, 640] 推理完成,输出数量: 1 输出形状: (1, 8400, 84)✅ 若能成功打印输出形状,则说明 ONNX 模型有效。
3. ONNX模型优化技巧
虽然 Ultralytics 默认导出的 ONNX 模型已具备基本可用性,但为进一步提升推理效率,建议进行以下优化。
### 3.1 使用 onnxsim 进行模型简化
即使设置了simplify=True,有时仍需手动运行简化工具以彻底清除冗余节点。
pip install onnxsim python -m onnxsim yolo26n.onnx yolo26n-sim.onnx该命令会生成更轻量、结构更清晰的yolo26n-sim.onnx,适用于 TensorRT 或 OpenVINO 等后端。
### 3.2 修改输入/输出节点名称(可选)
某些推理框架对输入名有特定要求(如必须为"input")。可通过修改 ONNX 图结构实现重命名。
示例代码片段(使用 onnx 库):
import onnx # 加载模型 model = onnx.load('yolo26n.onnx') # 修改输入名 model.graph.input[0].name = 'input' # 保存修改后的模型 onnx.save(model, 'yolo26n-renamed.onnx')💡 此操作常用于适配 TensorRT 的严格命名规则。
### 3.3 动态维度设置说明
当dynamic=True时,ONNX 模型支持灵活的输入尺寸。其输入定义如下:
input: float[batch_size, 3, height, width]可在推理时指定任意height和width(需为32倍数),例如(1, 3, 320, 320)或(4, 3, 736, 1280)。
若需固定输入尺寸(如嵌入式部署),可在导出时设置dynamic=False。
4. 常见问题与解决方案
### 4.1 导出失败:缺少 onnx 或 onnxsim 包
错误提示:
ModuleNotFoundError: No module named 'onnx'解决方法:
pip install onnx onnxsim### 4.2 推理报错:Unsupported ONNX operator GatherND
原因分析: 部分旧版本 ONNX Runtime 不支持 YOLO 中使用的高级索引操作(如GatherND)。
解决方案:
- 升级 ONNX Runtime 至最新版本:
pip install --upgrade onnxruntime-gpu - 或尝试降低
opset版本(如设为 11)再导出。
### 4.3 输出节点无法解析(shape mismatch)
现象: 输出 shape 为[1, 8400, 84],但不知如何解码为边界框。
说明: 这是 YOLOv8/v26 的标准输出格式:
8400是 anchor 数量(特征图总点数)84= 4(bbox)+ 80(COCO类别)
解码建议: 使用非极大值抑制(NMS)后处理,参考 Ultralytics 官方postprocess()实现逻辑,或借助onnxruntime-yolov8工具包自动解析。
### 4.4 GPU 导出失败:CUDA out of memory
解决策略:
- 减小
imgsz(如改为 320) - 设置
device='cpu'在 CPU 上导出 - 关闭
simplify后再单独运行简化步骤
5. 总结
5. 总结
本文围绕YOLO26 如何导出 ONNX 模型展开,结合官方训练与推理镜像环境,系统地介绍了从模型导出、格式验证到性能优化的全流程:
- 一键导出:利用
model.export(format='onnx')接口,轻松完成格式转换; - 动态输入支持:通过
dynamic=True实现多尺寸输入适配; - 模型简化:启用
simplify=True或手动调用onnxsim提升推理效率; - 推理验证:使用 ONNX Runtime 成功加载并运行推理,确认模型完整性;
- 常见问题应对:涵盖依赖缺失、算子不兼容、内存溢出等典型场景。
最终生成的 ONNX 模型可用于部署至 ONNX Runtime、TensorRT、OpenVINO、NCNN 等主流推理引擎,打通从训练到生产的“最后一公里”。
获取更多AI镜像
想探索更多AI镜像和应用场景?访问 CSDN星图镜像广场,提供丰富的预置镜像,覆盖大模型推理、图像生成、视频生成、模型微调等多个领域,支持一键部署。