YOLOv5模型导出为ONNX格式的全过程解析

YOLOv5 是一种高效的实时目标检测算法，而将 YOLOv5 模型导出为 ONNX 格式可以使其在不同的推理框架中使用，比如 TensorRT、ONNX Runtime 等。下面我们将详细解析 YOLOv5 模型导出为 ONNX 格式的全过程。

1. 环境准备

在开始导出模型之前，确保安装了以下依赖：

• Python >= 3.6
• PyTorch >= 1.7（建议使用与训练模型时一致的版本）
• ONNX >= 1.9
• YOLOv5 官方代码库

可以通过以下命令安装必要的依赖项：

pip install torch torchvision onnx

克隆 YOLOv5 的官方仓库并进入项目目录：

git clone https://github.com/ultralytics/yolov5.git
cd yolov5

2. 加载预训练模型

YOLOv5 提供了多种预训练模型，包括 yolov5s、yolov5m、yolov5l 和 yolov5x。我们以 yolov5s 为例，加载预训练权重文件。

import torch

# 加载预训练模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

# 设置模型为评估模式
model.eval()

如果需要加载自定义训练的模型，可以使用以下代码：

from models.experimental import attempt_load

# 加载自定义训练的权重文件
weights_path = 'path/to/custom_weights.pt'
model = attempt_load(weights_path, map_location=torch.device('cpu'))
model.eval()

3. 导出模型为 ONNX 格式

YOLOv5 提供了一个方便的脚本 export.py 来导出模型。以下是导出模型为 ONNX 格式的具体步骤。

3.1 使用 export.py 脚本

运行以下命令即可完成导出：

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

参数说明：

• --weights: 输入的 .pt 文件路径。
• --img: 图像输入尺寸，默认为 640x640。
• --batch: 批量大小，默认为 1。
• --include onnx: 指定导出为 ONNX 格式。

3.2 自定义导出过程

如果需要更灵活地控制导出过程，可以手动实现如下代码：

import torch.onnx

# 定义输入张量
dummy_input = torch.randn(1, 3, 640, 640)  # (batch_size, channels, height, width)

# 导出模型为 ONNX 格式
torch.onnx.export(
    model,
    dummy_input,
    "yolov5s.onnx",
    input_names=["input"],
    output_names=["output"],
    opset_version=11,
    do_constant_folding=True,
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

print("ONNX 模型导出成功！")

注意：

• opset_version 参数指定 ONNX 的操作集版本，推荐使用 11 或更高版本。
• dynamic_axes 参数允许输入和输出的批量大小动态变化。

4. 验证 ONNX 模型

导出完成后，可以使用 onnxruntime 验证模型是否正确。

import onnxruntime as ort
import numpy as np

# 加载 ONNX 模型
session = ort.InferenceSession("yolov5s.onnx")

# 准备输入数据
dummy_input = np.random.randn(1, 3, 640, 640).astype(np.float32)

# 运行推理
outputs = session.run(None, {"input": dummy_input})

print("ONNX 推理结果:", outputs)

如果推理结果与 PyTorch 模型一致，则说明导出成功。

5. 常见问题及解决方案

问题 1：导出失败，提示不支持的操作

某些自定义层可能无法直接转换为 ONNX 格式。解决方法是检查模型结构，替换或重写不支持的层。

问题 2：ONNX 模型过大

可以通过量化工具压缩模型大小。例如，使用 ONNX Runtime 的 Quantization 工具。

问题 3：动态输入导致错误

确保在导出时正确设置了 dynamic_axes 参数。

6. 总结

通过上述步骤，我们可以顺利将 YOLOv5 模型导出为 ONNX 格式，并验证其正确性。导出后的 ONNX 模型可以在多种硬件平台上高效运行，适用于嵌入式设备、服务器等场景。