YOLOv5模型转换实战：从pt到onnx的完整避坑指南（附常见错误排查）

YOLOv5模型转换实战：从PyTorch到ONNX的完整避坑指南

1. 模型转换基础准备

YOLOv5作为当前最流行的目标检测框架之一，其模型部署通常需要经过从PyTorch到ONNX的转换过程。在开始转换前，我们需要做好以下准备工作：

环境配置要求：

• Python 3.7或更高版本
• PyTorch 1.7+
• ONNX 1.8+
• ONNX Runtime 1.7+
• 最新版YOLOv5代码库

# 基础环境安装命令 pip install torch==1.10.0+cu113 torchvision==0.11.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html pip install onnx onnxruntime

文件结构检查： 确保YOLOv5项目目录包含以下关键文件：

yolov5/ ├── models/ │ ├── export.py # 转换脚本 │ ├── yolo.py # 模型定义 ├── weights/ │ └── best.pt # 训练好的模型权重

注意：建议使用官方推荐的PyTorch和ONNX版本组合，避免因版本不兼容导致的转换失败。

2. 基础转换流程详解

YOLOv5提供了原生的export.py脚本用于模型转换。以下是标准转换流程：

基本转换命令：

python models/export.py --weights yolov5s.pt --include onnx

关键参数解析：

输出验证： 成功转换后应生成以下文件：

• yolov5s.onnx(ONNX模型)
• yolov5s.torchscript.pt(可选)
• 文件大小应为原始PT文件的1.5-2倍

3. 常见问题与解决方案

3.1 路径错误与文件定位

典型报错：

FileNotFoundError: [Errno 2] No such file or directory: 'models/yolov5s.pt'

解决方案：

1. 确保从项目根目录运行脚本
2. 使用绝对路径更可靠：

python /path/to/yolov5/models/export.py --weights /path/to/weights.pt

3.2 输出文件异常

问题现象：

• 生成的ONNX文件大小异常（过小或过大）
• 推理结果不正确

排查步骤：

1. 检查转换日志是否有警告
2. 使用Netron可视化模型结构
3. 验证基础推理功能：

import onnxruntime import numpy as np ort_session = onnxruntime.InferenceSession("yolov5s.onnx") outputs = ort_session.run(None, {"images": np.random.randn(1,3,640,640).astype(np.float32)}) print(outputs[0].shape) # 应输出(1,25200,85)

3.3 动态维度处理

动态输入配置：

python export.py --weights yolov5s.pt --dynamic --batch-size 1

动态轴支持情况：

• 批处理维度（适合可变批处理）
• 图像尺寸（适合多分辨率输入）
• 输出检测数（适合不同目标数量）

提示：动态模型会增加部署复杂度，静态模型通常性能更优

4. 高级转换技巧

4.1 自定义算子处理

当遇到不支持的PyTorch算子时：

1. opset版本调整：

--opset 12 # 尝试更高版本

2. 自定义符号函数： 在export.py中添加：

torch.onnx.register_custom_op_symbolic( 'custom_op', custom_op_symbolic, opset_version)

4.2 模型优化技巧

优化前：

python -m onnxruntime.tools.convert_onnx_models_to_ort yolov5s.onnx

优化后对比：

4.3 多平台兼容性处理

平台特定建议：

• TensorRT：添加--end2end参数
• OpenVINO：使用mo.py二次转换
• CoreML：直接导出为CoreML格式

# TensorRT优化示例 trt_logger = trt.Logger(trt.Logger.WARNING) with trt.Builder(trt_logger) as builder: network = builder.create_network() parser = trt.OnnxParser(network, trt_logger) with open("yolov5s.onnx", "rb") as model: parser.parse(model.read())

5. 部署验证流程

完整的部署验证应包括以下步骤：

1. 精度验证：

# 对比原始模型和ONNX模型输出 pt_output = pt_model(test_img) onnx_output = ort_session.run(None, {"images": test_img.numpy()}) np.testing.assert_allclose(pt_output, onnx_output, rtol=1e-3)

2. 性能测试：

# 使用ONNX Runtime基准测试工具 onnxruntime_perf_test -m yolov5s.onnx -i input.json

3. 内存分析：

from onnxruntime.tools import onnx_model_utils model_info = onnx_model_utils.get_model_info("yolov5s.onnx") print(f"模型内存占用: {model_info.total_param_size/1024/1024:.2f}MB")

6. 性能调优实战

6.1 静态形状优化

固定输入尺寸：

python export.py --weights yolov5s.pt --img 640 --batch 1

优势对比：

6.2 量化压缩实践

FP16量化：

from onnxruntime.quantization import quantize_dynamic quantize_dynamic( "yolov5s.onnx", "yolov5s_quant.onnx", weight_type=QuantType.FLOAT16 )

精度-速度权衡数据：

7. 跨框架兼容性处理

7.1 算子兼容性矩阵

7.2 自定义层实现

Focus层替代方案：

class Focus(nn.Module): def forward(self, x): # 替换为标准的卷积+下采样 return self.conv(torch.cat([ x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2] ], 1))

8. 生产环境最佳实践

1. 版本固化：

# 保存环境配置 pip freeze > requirements.txt

2. 自动化验证脚本：

def validate_conversion(pt_model, onnx_path): # 加载ONNX模型 ort_session = onnxruntime.InferenceSession(onnx_path) # 生成测试数据 dummy_input = torch.randn(1,3,640,640) # 运行推理 pt_out = pt_model(dummy_input) onnx_out = ort_session.run(None, {"images": dummy_input.numpy()}) # 结果对比 assert np.allclose(pt_out.detach().numpy(), onnx_out[0], atol=1e-3)

3. 监控指标：

在实际项目中，我们发现将YOLOv5模型转换为ONNX格式时，保持PyTorch和ONNX版本的一致性至关重要。曾经遇到因PyTorch 1.8与ONNX 1.7不兼容导致的SiLU激活函数转换失败问题，升级到PyTorch 1.10后得到解决。