3、PyTorch模型导出为ONNX格式：torch.onnx.export核心参数详解

模型训练完了，权重文件躺在硬盘里。但说实话，这离真正能用还差一步。你想想看，PyTorch 模型跑起来得依赖 Python 环境、还得装 CUDA、装 torchvision，这玩意儿怎么塞进 C++ 写的 OpenCV 项目里？

所以我们需要一个中间格式——ONNX。它就像模型界的通用语言。今天我就带你把这步走通，重点聊聊 torch.onnx.export 那几个绕不开的参数。

本章核心目标： 把 PyTorch 的 YOLO 模型导出为 ONNX，并且能控制它输出单尺度还是多尺度。

3.1 为什么非得是 ONNX？

我刚开始做部署的时候，也想过直接拿 PyTorch 的 JIT trace 一把梭。结果到了 OpenCV 的 dnn 模块，发现支持得七零八落。ONNX 的好处在于：

  • 生态广：OpenCV、ONNX Runtime、TensorRT 全都认它。
  • 计算图固定：导出后就是个静态图，推理速度快，没有动态图的 overhead。
  • 算子标准化：不同框架之间互通的基石。

说白了，ONNX 就是模型界的 PDF。你发给我一个 PDF，我不管用 Chrome 还是 WPS 都能打开。ONNX 也是这个道理。

3.2 torch.onnx.export 核心参数

这个函数看着参数多，其实真正需要你操心的就那几个。我列个表，你一眼就能看明白。

参数名 作用 我的建议
model 要导出的 PyTorch 模型 记得切到 eval() 模式
args 模型的输入张量（用于 trace） 形状要和实际推理时一致
f 输出 ONNX 文件路径 用 .onnx 后缀
opset_version ONNX 算子集版本 我一般用 11 或 12，兼容性好
dynamic_axes 指定哪些维度是动态的 多尺度导出必须配这个
input_names 输入节点名称 随便起，但后面部署时要对应
output_names 输出节点名称 同上

3.2.1 opset_version：选对版本少踩坑

这个参数我吃过亏。有一次我图新鲜用了 opset 15，结果导出的模型在 OpenCV 4.5 上死活跑不起来。后来查文档才发现，OpenCV 的 ONNX 解析器只支持到 opset 11。

我的建议：

  • 如果你部署的目标环境是 OpenCV 4.5.x，用 opset_version=11。
  • 如果是 ONNX Runtime，可以放宽到 12 或 13。
  • 别追新。算子版本越高，兼容性风险越大。
小技巧： 导出时加一句 verbose=True，能看到每个算子被映射成了哪个 ONNX 版本。调试时特别好用。

3.2.2 dynamic_axes：让模型支持任意尺寸

YOLO 模型有个特点：输入图片不一定是固定尺寸。你可能训练时用了 640x640，但实际部署时用户传进来的是 1280x720。这时候就需要 dynamic_axes 出场了。

它的本质是告诉 ONNX：这个维度是可变的。比如 batch size 和图像的高宽。

# 定义动态轴
dynamic_axes = {
    'input': {0: 'batch_size', 2: 'height', 3: 'width'},
    'output': {0: 'batch_size', 2: 'height', 3: 'width'}
}

# 导出
torch.onnx.export(
    model,
    dummy_input,
    'yolo_model.onnx',
    opset_version=11,
    dynamic_axes=dynamic_axes,
    input_names=['input'],
    output_names=['output']
)

注意看，dynamic_axes 是个字典。key 是节点名，value 是 {维度索引: 维度名称}。维度名称你可以随便起，但最好见名知意。

注意： 动态轴不是万能的。有些算子（比如 Resize、GridSample）在动态形状下可能导出失败。我遇到过 YOLOv5 的 Focus 层在动态高宽下导出报错，后来换成 nn.Conv2d 才解决。

3.3 导出单尺度模型

单尺度最简单。你训练时用什么尺寸，导出时就用什么尺寸。比如 YOLOv8 默认 640x640。

import torch
import torch.onnx

# 假设 model 是你的 YOLO 模型
model.eval()

# 固定输入尺寸：batch=1, channel=3, height=640, width=640
dummy_input = torch.randn(1, 3, 640, 640)

torch.onnx.export(
    model,
    dummy_input,
    'yolov8_single_scale.onnx',
    opset_version=11,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes=None  # 不设动态轴，就是固定尺寸
)

print("单尺度模型导出完成")

这种模型跑起来最快，因为所有算子都能做 shape 特化优化。缺点嘛，就是输入图片必须 resize 到 640x640，可能会变形。

3.4 导出多尺度模型

多尺度模型更灵活。你可以在推理时传入任意尺寸的图片，模型自动适配。这在实际项目中太常见了——用户上传的图片哪有那么规矩？

导出多尺度模型，关键就是配好 dynamic_axes。

import torch
import torch.onnx

model.eval()

# dummy 输入用任意尺寸，比如 640x640
dummy_input = torch.randn(1, 3, 640, 640)

# 定义动态轴：batch、height、width 都可变
dynamic_axes = {
    'input': {0: 'batch_size', 2: 'height', 3: 'width'},
    'output': {0: 'batch_size', 2: 'height', 3: 'width'}
}

torch.onnx.export(
    model,
    dummy_input,
    'yolov8_multi_scale.onnx',
    opset_version=11,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes=dynamic_axes
)

print("多尺度模型导出完成")

导出后，你用 onnxruntime 或者 OpenCV 加载时，可以传入 (1, 3, 480, 640) 或者 (1, 3, 720, 1280)，模型都能正常推理。

经验之谈： 多尺度模型在 OpenCV 的 dnn 模块里，对 blobFromImage 的尺寸参数要求更宽松。你传进去的尺寸只要和 ONNX 定义的动态轴兼容就行，不用非得是 32 的倍数（但 YOLO 内部有下采样，最好还是保持 32 对齐）。

3.5 单尺度 vs 多尺度：怎么选？

我直接给你一个决策表：

场景 推荐方案 原因
固定摄像头分辨率（如 1080p） 单尺度 速度快，无动态 shape 开销
用户上传图片，尺寸不固定 多尺度 灵活，不用每次都 resize 到固定尺寸
需要 batch 推理（多张图一起） 多尺度（batch 维度动态） batch size 可变，吞吐量更高
部署到嵌入式设备（如 Jetson） 单尺度 内存占用固定，算子优化更彻底

3.6 导出后的验证

模型导出来，别急着部署。先验证一下 ONNX 文件对不对。我习惯用 onnxruntime 跑一遍推理，和 PyTorch 的结果对比。

import onnxruntime as ort
import numpy as np

# 加载 ONNX 模型
session = ort.InferenceSession('yolov8_multi_scale.onnx')

# 构造输入（注意：ONNX 的输入是 numpy 数组）
input_name = session.get_inputs()[0].name
dummy_input_np = np.random.randn(1, 3, 640, 640).astype(np.float32)

# 推理
outputs = session.run(None, {input_name: dummy_input_np})

print("ONNX 推理成功，输出形状：", outputs[0].shape)

如果这一步没报错，恭喜你，模型导出成功了。接下来就可以扔给 OpenCV 了。

一个小坑： 有些 YOLO 模型在导出时，后处理（NMS）也会被包含在计算图里。我个人建议把 NMS 留在外面做，ONNX 只负责输出原始检测结果（比如 [batch, num_boxes, 6]）。这样部署时更灵活。

3.7 本章知识结构图

下面这张图帮你理清整个导出流程和关键决策点：

PyTorch → ONNX 导出流程 PyTorch 模型 选择导出策略 单尺度导出 dynamic_axes=None 多尺度导出 dynamic_axes 配置 opset_version=11 opset_version + dynamic_axes 单尺度 ONNX 多尺度 ONNX 💡 关键参数 opset_version: 兼容性 dynamic_axes: 灵活性

嗯，这张图把整个决策路径画清楚了。你从 PyTorch 模型出发，先决定要单尺度还是多尺度，然后配好对应的参数，最后导出 ONNX。每一步都有对应的参数在控制。

好了，这一章的内容就到这儿。你手头如果有训练好的 YOLO 模型，现在就可以试着导出一版。记住：先单尺度跑通，再上多尺度。别一上来就搞复杂的，容易把自己绕进去。

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