YOLOv9模型导出ONNX?后续推理格式转换路径
1. 镜像环境说明
- 核心框架: pytorch==1.10.0
- CUDA版本: 12.1
- Python版本: 3.8.5
- 主要依赖: torchvision==0.11.0,torchaudio==0.10.0,cudatoolkit=11.3, numpy, opencv-python, pandas, matplotlib, tqdm, seaborn等。
- 代码位置:
/root/yolov9
本镜像基于 YOLOv9 官方代码库构建,预装了完整的深度学习开发环境,集成了训练、推理及评估所需的所有依赖,开箱即用。用户无需手动配置复杂环境即可快速开展目标检测任务的开发与部署工作。
2. 快速上手
2.1 激活环境
在使用前,请确保激活对应的 Conda 环境:
conda activate yolov92.2 模型推理 (Inference)
进入 YOLOv9 项目目录:
cd /root/yolov9运行以下命令进行图像检测:
python detect_dual.py --source './data/images/horses.jpg' --img 640 --device 0 --weights './yolov9-s.pt' --name yolov9_s_640_detect检测结果将保存在runs/detect/yolov9_s_640_detect目录中,包含标注框和类别信息的可视化图像。
2.3 模型训练 (Training)
单卡训练示例命令如下:
python train_dual.py --workers 8 --device 0 --batch 64 --data data.yaml --img 640 --cfg models/detect/yolov9-s.yaml --weights '' --name yolov9-s --hyp hyp.scratch-high.yaml --min-items 0 --epochs 20 --close-mosaic 15该命令使用自定义数据集(需提前准备data.yaml),以yolov9-s架构从零开始训练 20 轮。
3. 已包含权重文件
镜像内已预下载轻量级模型权重yolov9-s.pt,位于/root/yolov9根目录下,可直接用于推理或微调任务,避免重复下载带来的网络问题。
4. 常见问题
- 数据集准备:请将您的数据集按照 YOLO 格式组织(每张图对应一个
.txt标注文件),并在data.yaml中正确设置train,val,nc,names字段。 - 环境激活失败:镜像启动后默认处于
base环境,务必执行conda activate yolov9切换至目标环境后再运行脚本。 - 显存不足:若出现 CUDA out of memory 错误,建议降低
--batch批大小或调整输入分辨率--img。
5. YOLOv9 模型导出 ONNX 流程详解
YOLOv9 支持通过官方脚本导出为 ONNX(Open Neural Network Exchange)格式,便于跨平台部署与推理引擎集成(如 TensorRT、ONNX Runtime、OpenVINO 等)。以下是详细操作步骤。
5.1 导出命令说明
YOLOv9 提供了统一的导出接口export.py,支持多种格式输出。要导出为 ONNX 格式,执行以下命令:
python export.py --weights ./yolov9-s.pt --include onnx --imgsz 640 --batch 1 --dynamic参数解释:
--weights: 指定待导出的 PyTorch 权重路径。--include onnx: 指定导出格式为 ONNX。--imgsz 640: 输入尺寸,必须与推理时一致。--batch 1: 推理批次大小,设为 1 可保证兼容性。--dynamic: 启用动态轴(dynamic axes),允许变尺寸输入,适用于不同分辨率图像。
注意:若不加
--dynamic,生成的 ONNX 模型将固定输入维度,限制部署灵活性。
5.2 输出文件分析
执行成功后,会在当前目录生成名为yolov9-s.onnx的文件。可通过 Netron 工具加载查看网络结构,确认输入输出节点名称及维度。
典型输出结构如下: -输入节点名:images-输入形状:(1, 3, 640, 640)(当未启用 dynamic 时) -输出节点名:output或类似命名(具体取决于实现)
5.3 验证 ONNX 模型有效性
建议使用onnxruntime对导出模型进行前向推理验证,确保数值一致性。
安装 ONNX Runtime:
pip install onnxruntime-gpuPython 验证代码片段:
import onnxruntime as ort import cv2 import numpy as np # 加载 ONNX 模型 session = ort.InferenceSession("yolov9-s.onnx", providers=["CUDAExecutionProvider"]) # 图像预处理 image = cv2.imread('./data/images/horses.jpg') image = cv2.resize(image, (640, 640)) image = image.transpose(2, 0, 1).astype(np.float32) / 255.0 image = np.expand_dims(image, axis=0) # 推理 outputs = session.run(None, {"images": image}) print("ONNX Inference Output Shape:", outputs[0].shape)若能正常输出(1, N, 85)形状的结果(N 为候选框数量,85 = 4 bbox + 1 confidence + 80 class scores),则表示导出成功。
6. 后续推理格式转换路径
ONNX 是模型中间表示的关键桥梁,支持进一步转换为多种高性能推理格式。以下是常见部署路径及其工具链。
6.1 转换为 TensorRT 引擎
TensorRT 可显著提升 NVIDIA GPU 上的推理速度,尤其适合边缘设备(如 Jetson 系列)或服务器端高并发场景。
转换流程:
- 使用
trtexec工具直接编译 ONNX 至 TRT 引擎:bash trtexec --onnx=yolov9-s.onnx --saveEngine=yolov9-s.engine --fp16 --workspace=4096 - 在 C++ 或 Python 中加载
.engine文件进行高速推理。
提示:若遇到算子不支持问题,可尝试简化 ONNX 模型或升级 TensorRT 版本(推荐 8.6+)。
6.2 转换为 OpenVINO IR 模型
适用于 Intel CPU/GPU/VPU 设备(如 Movidius Myriad X)。
转换步骤:
- 安装 OpenVINO Toolkit:
bash pip install openvino-dev - 使用
mo命令行工具转换:bash mo --input_model yolov9-s.onnx --output_dir ir_model --data_type FP16
生成的.xml和.bin文件可在 OpenVINO Runtime 中高效加载。
6.3 使用 ONNX Runtime 直接部署
对于通用 CPU/GPU 推理需求,ONNX Runtime 是最轻量且跨平台的选择。
部署优势:
- 支持 Windows/Linux/macOS/Android/iOS
- 多执行后端(CPU、CUDA、Core ML、DML)
- 易于集成到 Web 服务(Flask/FastAPI)
示例部署架构:
from fastapi import FastAPI import uvicorn app = FastAPI() model = ort.InferenceSession("yolov9-s.onnx", providers=["CUDAExecutionProvider"]) @app.post("/detect") def detect_image(image: UploadFile): # 预处理 → 推理 → 后处理 → 返回 JSON 结果 pass if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=8000)7. 总结
本文围绕 YOLOv9 官方训练与推理镜像展开,系统介绍了其环境配置、快速上手方法,并重点讲解了如何将训练好的模型导出为 ONNX 格式。在此基础上,进一步梳理了主流的后续推理格式转换路径,包括:
- TensorRT:适用于 NVIDIA GPU 的极致性能优化
- OpenVINO:面向 Intel 硬件的高效推理方案
- ONNX Runtime:跨平台、易集成的通用部署选择
通过 ONNX 这一标准化中间层,开发者可以灵活适配不同硬件平台,实现“一次训练,多端部署”的工程目标。结合本镜像提供的完整依赖环境,用户可专注于模型优化与业务集成,大幅提升开发效率。
获取更多AI镜像
想探索更多AI镜像和应用场景?访问 CSDN星图镜像广场,提供丰富的预置镜像,覆盖大模型推理、图像生成、视频生成、模型微调等多个领域,支持一键部署。