如何将自定义YoloV5 PyTorch模型转换为适配OpenCV的格式？

我正尝试通过ONNX将自定义YoloV5s PyTorch模型转换为可用于OpenCV的格式，但出现了一些错误，希望有人能提供帮助。以下是我的代码：

import torch
import onnx
from onnxsim import simplify


# Step 1: Load the YOLOv5 model
model = torch.hub.load('https://github.com/ultralytics/yolov5','custom',
                       'best.pt')
batch_size = 16
# Step 2: Export the model to ONNX format
torch.onnx.export(model,
                  torch.zeros(batch_size, 3, 640, 640),
                  "yolov5s.onnx",
                  opset_version=11)

# Step 3: Simplify the exported ONNX model
onnx_model = onnx.load("yolov5s.onnx")
simplified_model, _ = simplify(onnx_model)
onnx.save(simplified_model, "yolov5s.onnx")

问题排查与修正方案 #

• 版本兼容调整：YOLOv5不同版本的ONNX导出逻辑差异较大，建议使用v6.x或v7.x的稳定版本，避免版本不匹配导致的导出异常。
• 适配OpenCV的导出参数：OpenCV的dnn模块对ONNX模型有特定要求，调整导出参数：
  • 将opset_version提升至12（OpenCV 4.5+支持该版本，兼容性更好）
  • 添加do_constant_folding=True优化常量折叠逻辑
  • 指定input_names和output_names，方便OpenCV识别输入输出节点
  • 若需要动态batch，添加dynamic_axes参数；否则固定batch尺寸即可
• 模型有效性检查：简化前先通过onnx.checker.check_model()验证原始ONNX模型是否合法，避免简化过程中出现错误。
• 简化时指定输入形状：给onnxsim传入input_shapes参数，帮助工具更精准地优化模型结构。

修改后的示例代码 #

import torch
import onnx
from onnxsim import simplify

# 加载自定义YOLOv5模型并设置为评估模式
model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')
model.eval()
batch_size = 16
input_shape = (batch_size, 3, 640, 640)

# 导出适配OpenCV的ONNX模型
torch.onnx.export(
    model,
    torch.zeros(input_shape),
    "yolov5s_opencv.onnx",
    opset_version=12,
    do_constant_folding=True,
    input_names=['images'],
    output_names=['output'],
    # 如需支持动态batch，取消下面注释
    # dynamic_axes={'images': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)

# 检查并简化模型
onnx_model = onnx.load("yolov5s_opencv.onnx")
onnx.checker.check_model(onnx_model)
simplified_model, check = simplify(onnx_model, input_shapes={'images': input_shape})
assert check, "简化后的ONNX模型验证失败"
onnx.save(simplified_model, "yolov5s_opencv_simplified.onnx")

OpenCV加载验证示例 #

确保OpenCV版本≥4.4，预处理逻辑与YOLOv5保持一致：

import cv2
import numpy as np

# 加载模型
net = cv2.dnn.readNetFromONNX("yolov5s_opencv_simplified.onnx")
# 预处理图像
img = cv2.imread("test.jpg")
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_resized = cv2.resize(img_rgb, (640, 640))
img_normalized = img_resized / 255.0
input_blob = np.expand_dims(img_normalized.transpose(2, 0, 1), axis=0).astype(np.float32)
# 执行推理
net.setInput(input_blob)
output = net.forward()

内容的提问来源于stack exchange，提问作者invisible