模型获取与转换：从PyTorch/TensorFlow导出ONNX模型
核心原理：为什么需要“导出”而不是“保存”
在PyTorch或TensorFlow中训练得到的模型，本质上是与特定框架运行时深度绑定的计算图与权重集合。PyTorch的.pth文件保存了Python对象的序列化状态，TensorFlow的.pb或.h5则依赖其自身的图执行引擎。而ONNX（Open Neural Network Exchange）定义了一套与框架无关的静态计算图协议，它将算子（Operator）、张量形状、数据流方向统一为中间表示（IR）。导出ONNX的过程，实质上是将框架动态图“追踪”或“转换”为静态图，并逐一映射到ONNX算子集的过程。
这一过程的关键难点在于：动态控制流（如if、for依赖数据）无法被静态图直接表达，自定义算子需要注册映射，张量形状在动态维度下需要显式声明。理解这些底层机制，才能避免“导出成功但推理报错”的陷阱。
核心概念速览
- 追踪导出（Tracing）：PyTorch默认方式，用示例输入执行模型，记录实际执行路径，速度快但无法捕获未执行的分支。
- 脚本化（Scripting）：解析Python源码为TorchScript，支持动态控制流，但需要模型代码符合脚本化约束。
- 静态图转换：TensorFlow通过
tf2onnx将冻结的GraphDef转换为ONNX，需先冻结变量为常量。 - 算子映射表：框架算子与ONNX算子的对应关系，映射缺失时需自定义转换或分解。
方法流程：PyTorch导出ONNX的标准步骤
PyTorch自1.2版本起，torch.onnx.export成为官方推荐的导出接口。其核心参数包括model（继承自nn.Module的实例）、args（示例输入张量）、f（输出文件名）、input_names与output_names（用于指定IO别名），以及dynamic_axes（声明动态维度）。以下是一个典型的导出流程：
import torch
import torchvision.models as models
# 1. 加载模型并设为评估模式
model = models.resnet18(pretrained=True)
model.eval()
# 2. 构造示例输入（batch=1, 3通道, 224x224）
dummy_input = torch.randn(1, 3, 224, 224)
# 3. 导出ONNX
torch.onnx.export(
model, # 模型
dummy_input, # 示例输入
"resnet18.onnx", # 输出路径
export_params=True, # 导出权重
opset_version=17, # ONNX算子集版本
do_constant_folding=True, # 常量折叠优化
input_names=["input"], # 输入名
output_names=["output"], # 输出名
dynamic_axes={
"input": {0: "batch_size"}, # 第0维动态
"output": {0: "batch_size"}
}
)
print("导出完成")导出后，务必使用onnx.checker验证模型结构完整性，再用onnxruntime进行数值一致性比对。若输出误差超过1e-4，需检查是否有不支持的算子或精度问题。
动态维度声明技巧：若模型需要处理可变batch或可变分辨率，必须在dynamic_axes中显式声明。否则ONNX会将示例输入的维度固化为静态值，推理时遇到不同形状会直接报错。推荐将batch维设为动态，而H/W维尽量固定，以减少算子优化难度。
方法流程：TensorFlow模型转换ONNX
TensorFlow模型转换通常借助tf2onnx工具链。其基本流程分为三步：导出SavedModel → 冻结图（将变量转为常量） → 调用tf2onnx转换。对于Keras模型，可直接使用tf2onnx.convert.from_keras简化流程。
import tensorflow as tf
import tf2onnx
# 方式一：从Keras模型直接转换
model = tf.keras.applications.MobileNetV2(weights="imagenet")
model.save("mobilenetv2_savedmodel")
# 使用tf2onnx命令行转换
# python -m tf2onnx.convert \
# --saved-model mobilenetv2_savedmodel \
# --output mobilenetv2.onnx \
# --opset 17 \
# --inputs-as-nchw input_1:0
# 方式二：Python API转换
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
output_path = tf2onnx.convert.from_keras(model, input_signature=spec, opset=17)
# 输出为 (onnx_graph, success_status)TensorFlow的NHWC内存布局与ONNX默认的NCHW不同，转换时需通过--inputs-as-nchw参数指定输入通道顺序，否则推理性能会大幅下降。此外，TF的某些复合算子（如FusedBatchNorm）在转换时可能被拆分为多个基础算子，导致图规模膨胀，建议在转换前使用graph_transform工具进行优化。
案例：YOLOv5从PyTorch到ONNX的实战导出
目标检测模型YOLOv5的导出是典型场景。其模型包含多输出头（边界框回归、类别预测、对象置信度）和非极大值抑制（NMS）后处理。由于ONNX标准算子集中不包含NMS，导出时需将后处理逻辑剥离，仅保留网络前向部分。
# YOLOv5官方导出脚本核心命令
python export.py --weights yolov5s.pt --include onnx --opset 17 --dynamic
# 导出后使用onnxruntime验证
import onnxruntime as ort
import numpy as np
sess = ort.InferenceSession("yolov5s.onnx")
input_name = sess.get_inputs()[0].name
output_names = [o.name for o in sess.get_outputs()]
# 模拟输入 (batch=1, 3, 640, 640)
img = np.random.randn(1, 3, 640, 640).astype(np.float32)
outputs = sess.run(output_names, {input_name: img})
print(f"输出张量数量: {len(outputs)}")
for i, out in enumerate(outputs):
print(f"输出{i}形状: {out.shape}")实际部署时，NMS通常由ONNX Runtime的NonMaxSuppression算子或TensorRT的插件层实现。若目标平台不支持，则需在C++端自行实现NMS逻辑。导出时建议将输出头分离为多个独立输出，便于后续灵活拼接后处理流程。
误区与复盘：常见导出失败原因及对策
| 常见错误 | 根因分析 | 解决方案 |
|---|---|---|
| 导出时报“Unsupported operator” | 模型使用了ONNX算子集之外的算子（如自定义RoIAlign） | 将自定义算子分解为ONNX基础算子，或注册自定义算子实现 |
| 推理结果与PyTorch差异大 | 模型处于训练模式（Dropout/BatchNorm行为不同） | 导出前必须调用model.eval() |
| 动态batch推理报错 | 未在dynamic_axes中声明batch维度 | 重新导出并指定dynamic_axes={"input": {0: "batch"}} |
| TensorFlow转换后精度下降 | NHWC与NCHW布局未正确转换 | 使用--inputs-as-nchw参数，并在验证时对比中间层输出 |
| ONNX文件过大（>500MB） | 权重未压缩或包含冗余计算图 | 使用onnx-simplifier简化图结构，或对权重进行量化 |
关键复盘提醒：导出ONNX并非“一键完成”的简单操作。建议在导出后立即执行三项验证：① onnx.checker.check_model检查结构合法性；② 使用onnxruntime跑一遍随机输入，确认无运行时错误；③ 对比PyTorch/TensorFlow与ONNX的输出张量，确保数值误差在可接受范围（通常<1e-4）。任何一步失败，都应回溯至模型定义或导出参数，而不是盲目修改推理代码。
掌握从PyTorch和TensorFlow导出ONNX的核心方法，是打通训练框架与C++推理引擎的关键桥梁。下一阶段，我们将深入探讨ONNX模型的结构解析与算子兼容性分析，为后续的优化与部署打下坚实基础。