2. YOLO模型导出为ONNX

把训练好的YOLO模型转成ONNX，是整个部署流程的第一步。这一步要是没走稳，后面全是坑。我个人习惯把ONNX看作“模型界的通用语言”——不管你是PyTorch、TensorFlow还是别的框架训练出来的，ONNX都能把它们统一起来。

说白了，ONNX就是一个中间表示。它不依赖任何训练框架，只描述计算图的结构和权重。这样，我们就能用统一的推理引擎去跑模型了。

本章核心目标： 掌握从PyTorch导出YOLO模型到ONNX的完整流程，包括动态轴设置、opset版本选择，以及导出后的验证方法。

2.1 PyTorch模型导出基础

PyTorch官方提供了torch.onnx.export这个接口。嗯，这里要注意，导出不是直接把模型扔进去就完事了。你需要告诉它三件事：模型本身、一个示例输入、以及一些配置参数。

先看一个最基本的导出代码：

import torch
import torch.onnx

# 加载你的YOLO模型
model = torch.load('yolov8n.pt', map_location='cpu')
model.eval()

# 创建一个示例输入（batch_size=1, 3通道, 640x640）
dummy_input = torch.randn(1, 3, 640, 640)

# 导出ONNX
torch.onnx.export(
    model,                    # 模型
    dummy_input,              # 示例输入
    'yolov8n.onnx',           # 输出文件名
    export_params=True,       # 导出权重参数
    opset_version=11,         # opset版本
    do_constant_folding=True, # 常量折叠优化
    input_names=['input'],    # 输入名称
    output_names=['output'],  # 输出名称
    dynamic_axes={}           # 动态轴设置
)

这段代码看起来简单，但我在项目中遇到过一个问题：如果不设置do_constant_folding=True，导出的模型里会残留一些冗余的计算节点，推理速度会慢10%-15%。所以这个参数我建议一直开着。

2.2 动态轴设置

为什么要设置动态轴？你想想看，训练的时候batch size可能是固定的，但部署的时候，你不可能每次都只处理一张图。有时候一次来10张，有时候一次来100张。如果ONNX模型把batch size写死了，那就只能重新导出。

动态轴设置就是告诉ONNX：这个维度是可变的。来看具体怎么配：

dynamic_axes = {
    'input': {
        0: 'batch_size',  # 动态batch
        2: 'height',      # 动态高度
        3: 'width'        # 动态宽度
    },
    'output': {
        0: 'batch_size'   # 输出也要对应动态
    }
}

torch.onnx.export(
    model,
    dummy_input,
    'yolov8n_dynamic.onnx',
    export_params=True,
    opset_version=11,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes=dynamic_axes  # 传入动态轴配置
)
我的经验： 动态轴不是越多越好。如果你确定部署时输入尺寸固定（比如永远是640x640），那就只设batch为动态。设太多动态轴会增加推理引擎的优化难度，反而可能降低性能。

这里有个细节：YOLO模型的输出通常包含多个尺度的检测结果。如果你把height和width也设为动态，那输出层的维度也会跟着变。推理引擎需要额外处理这种动态形状，有些算子可能不支持。所以，我一般只动态batch size，其他维度固定。

2.3 opset版本选择

opset版本，说白了就是ONNX的“方言版本”。不同版本支持的算子不一样。选错了，导出的时候可能报错，或者导出来跑不了。

我整理了一个常用opset版本的对比表：

opset版本 支持的关键算子 适用场景 我的建议
10 基础卷积、池化、激活 老模型、老推理引擎 不推荐，算子太少
11 Resize、Slice、Scatter YOLOv5/v8常见选择 最稳妥，兼容性好
12 ReduceL2、SoftmaxCrossEntropy 需要新算子的模型 可以尝试
13 Softmax改进、Gelu Transformer类模型 YOLO用不到
15 Trilu、Cumsum 最新模型 兼容性风险高
避坑指南： 我曾经选过opset=15导出YOLOv8，结果在TensorRT上跑的时候，有个Slice算子不支持，折腾了两天才找到替代方案。所以，除非你的推理引擎明确支持高版本opset，否则我建议用opset=11，这是最安全的选择。

为什么opset=11最稳妥？因为大多数推理引擎（ONNX Runtime、TensorRT、OpenVINO）都对opset=11做了充分优化。你想想看，如果选了一个太新的版本，推理引擎还没来得及适配，那就只能自己手写算子插件了——这工作量可不小。

2.4 验证导出的ONNX模型

导出完了不等于万事大吉。我见过太多人导出ONNX后直接扔给推理引擎，结果跑出来一堆乱码。验证这一步，不能省。

验证分三步走：

  1. 结构验证： 用ONNX官方工具检查模型结构是否完整
  2. 数值验证： 对比PyTorch和ONNX的输出是否一致
  3. 形状验证： 检查输入输出形状是否符合预期

先看结构验证的代码：

import onnx

# 加载ONNX模型
onnx_model = onnx.load('yolov8n.onnx')

# 检查模型结构
try:
    onnx.checker.check_model(onnx_model)
    print("✅ 模型结构验证通过")
except onnx.checker.ValidationError as e:
    print(f"❌ 模型结构验证失败: {e}")

这一步能检查出模型里有没有无效的节点、缺失的权重、或者不兼容的算子。如果报错了，别慌，先看看是不是opset版本选错了。

接下来是数值验证——这一步最关键：

import onnxruntime as ort
import numpy as np

# 创建ONNX Runtime推理会话
ort_session = ort.InferenceSession('yolov8n.onnx')

# 准备输入数据（用和导出时一样的随机数）
test_input = np.random.randn(1, 3, 640, 640).astype(np.float32)

# PyTorch推理
with torch.no_grad():
    torch_output = model(torch.from_numpy(test_input)).numpy()

# ONNX Runtime推理
ort_inputs = {ort_session.get_inputs()[0].name: test_input}
ort_output = ort_session.run(None, ort_inputs)[0]

# 对比输出
diff = np.abs(torch_output - ort_output).max()
print(f"最大数值差异: {diff:.6f}")

if diff < 1e-3:
    print("✅ 数值验证通过")
else:
    print("❌ 数值差异过大，请检查导出过程")
关键指标： 最大数值差异小于1e-3就算通过。如果差异在1e-2级别，说明模型精度已经受损，需要检查是否有算子被错误替换。

最后是形状验证：

# 检查输入输出形状
input_shape = ort_session.get_inputs()[0].shape
output_shape = ort_session.get_outputs()[0].shape

print(f"输入形状: {input_shape}")
print(f"输出形状: {output_shape}")

# 检查动态轴是否生效
# 如果设置了动态batch，这里应该显示为 'batch_size' 或 -1
assert input_shape[0] == 'batch_size' or input_shape[0] == -1, "动态batch未生效"

嗯，这里要注意：如果动态轴设置正确，input_shape[0]会显示为'batch_size'（字符串）或者-1。如果显示的是1，说明动态轴没生效，回去检查dynamic_axes的配置。

2.5 本章知识体系

我把整个导出流程画成了一张图，方便你理解各个环节的关系：

YOLO模型导出ONNX全流程 PyTorch训练模型 配置导出参数 动态轴设置 opset版本选择 输入输出命名 导出ONNX模型 三步验证：结构→数值→形状 💡 验证通过后才能进入下一步：模型量化与推理引擎适配
我的工作流： 每次导出ONNX后，我会写一个自动化脚本，把结构验证、数值验证、形状验证串起来跑一遍。只要有一个验证失败，脚本就自动发邮件告警。这样能避免把有问题的模型推到生产环境。

好了，这一章的内容就这些。记住三个关键点：动态轴只设batch、opset选11、验证三步不能省。把这些基础打牢了，后面的量化、剪枝、推理引擎适配才能顺利推进。

← 上一章 📖 返回目录 下一章 →