第4章 模型转换与中间表示：PyTorch/TF模型导出为ONNX
模型训练好了，怎么塞进嵌入式设备？
直接拿PyTorch的.pt文件去跑？别想了。嵌入式推理引擎不认识这些框架的私有格式。它们只认一个东西——中间表示。而ONNX，就是目前最通用的那个中间人。
我个人习惯把ONNX比作「模型界的通用语言」。你拿PyTorch写个网络，我用TensorFlow搭个结构，最后都转成ONNX，大家就能在同一个平台上对话了。这一章，我们就来聊聊怎么把训练好的模型「翻译」成ONNX，以及翻译过程中容易踩的坑。
4.1 为什么要用ONNX？
说白了，ONNX解决的是「框架锁死」的问题。你想想看，项目做到一半，发现某个算子PyTorch支持但TensorRT不支持，或者反过来——这时候如果没有ONNX，你就得重写模型。我早期做项目时就被这种问题折腾过，后来学乖了，所有模型先过一遍ONNX再说。
| 特性 | 原生框架模型 | ONNX中间表示 |
|---|---|---|
| 跨框架兼容 | 仅限本框架 | PyTorch/TF/MXNet互通 |
| 推理后端支持 | 有限 | TensorRT/OpenVINO/NCNN等主流引擎均支持 |
| 算子标准化 | 框架自定义 | 统一算子集，便于优化 |
| 部署流程 | 需要框架运行时 | 轻量级runtime即可 |
核心观点：ONNX不是推理引擎，它是模型流通的「中间货币」。你最终还是要把它转成目标平台的具体格式（比如TensorRT的.engine或NCNN的.param）。但有了ONNX，你只需要做一次转换工作。
4.2 PyTorch模型导出为ONNX
PyTorch官方提供了torch.onnx.export接口。用法不复杂，但参数调不对，导出来的模型可能跑不动。
4.2.1 基本导出流程
import torch
import torch.onnx
# 假设你有一个训练好的模型
model = YourModel()
model.load_state_dict(torch.load('model.pth'))
model.eval()
# 构造一个 dummy 输入
dummy_input = torch.randn(1, 3, 224, 224)
# 导出 ONNX
torch.onnx.export(
model, # 模型
dummy_input, # 示例输入
'model.onnx', # 输出路径
export_params=True, # 是否导出参数
opset_version=11, # ONNX算子集版本
do_constant_folding=True, # 常量折叠优化
input_names=['input'], # 输入节点名称
output_names=['output'], # 输出节点名称
dynamic_axes={
'input': {0: 'batch_size'},
'output': {0: 'batch_size'}
}
)
嗯，这里要注意opset_version。我建议至少用11以上。版本太低，很多新算子不支持；版本太高，老推理引擎又认不出来。我在项目中遇到过用opset=13导出的模型，结果客户的NVIDIA驱动太老，TensorRT不认——最后只能降版本重导。
4.2.2 dynamic_axes 的坑
很多新手会忽略dynamic_axes。如果你不设置，导出的ONNX默认输入形状是固定的——比如你dummy_input是1x3x224x224，那模型就只能处理batch_size=1的输入。想跑batch_size=4？不行。
设置dynamic_axes后，ONNX会把这个维度标记为动态。但注意，不是所有算子都支持动态形状。我踩过的一个坑是：模型里有个reshape操作，动态batch下reshape的维度计算错了，推理时直接崩掉。
避坑指南：导出时设置dynamic_axes后，一定要用不同batch_size的输入做一次推理验证。我曾经因为偷懒没验证，结果上线后模型在batch_size=2时输出全为0，排查了两天才发现是某个自定义算子没处理好动态维度。
4.3 TensorFlow模型导出为ONNX
TF模型转ONNX，主流工具是tf2onnx。流程跟PyTorch类似，但细节上有些差异。
4.3.1 使用tf2onnx命令行工具
# 安装
pip install tf2onnx
# 转换
python -m tf2onnx.convert \
--saved-model ./saved_model \
--output model.onnx \
--opset 11 \
--inputs input:0 \
--outputs output:0
这里有个容易忽略的点：TF的模型输入输出名称。你用saved_model_cli show --dir ./saved_model --all先看一眼，确认输入输出节点的名字。我见过有人把名字写错，转换出来的ONNX输入输出全是乱序的。
4.3.2 TF特有算子的处理
TF有些算子ONNX原生不支持，比如FusedBatchNorm、TopKV2。tf2onnx内部会做算子映射，但映射不一定完美。我遇到过FusedBatchNorm转成ONNX后精度掉了0.5%，排查发现是epsilon参数没传对。
| TF算子 | ONNX映射算子 | 常见问题 |
|---|---|---|
| FusedBatchNorm | BatchNormalization | epsilon参数精度丢失 |
| TopKV2 | TopK | K值必须为常量 |
| ResizeBilinear | Resize | 坐标变换模式不匹配 |
| SplitV | Split | split维度索引差异 |
个人经验：TF转ONNX后，建议用onnxruntime跑一遍推理，跟原始TF模型的输出做数值对比。误差在1e-5以内算正常，超过1e-3就要警惕了。
4.4 ONNX结构解析
ONNX文件本质上是一个protobuf序列化的数据。你可以用onnx库来解析它。
4.4.1 查看模型结构
import onnx
model = onnx.load('model.onnx')
print(model.graph.node) # 打印所有算子节点
print(model.graph.input) # 输入信息
print(model.graph.output) # 输出信息
# 查看某个节点的属性
for node in model.graph.node:
if node.op_type == 'Conv':
print(node.name, node.attribute)
我个人习惯用netron这个可视化工具。它能把ONNX模型画成一张网络图，节点、连接、参数一目了然。你可以在浏览器里打开，拖拽查看。对于排查模型结构问题，比看代码快多了。
4.4.2 修改ONNX模型
有时候你需要手动修改ONNX模型——比如删除某个节点、修改输入形状。可以用onnx库的API来做。
import onnx
from onnx import helper, TensorProto
model = onnx.load('model.onnx')
# 修改输入形状
new_input = helper.make_tensor_value_info(
'input', TensorProto.FLOAT, [None, 3, 224, 224]
)
model.graph.input[0].CopyFrom(new_input)
# 删除某个节点
node_to_remove = 'node_name'
model.graph.node.remove(
[n for n in model.graph.node if n.name == node_to_remove][0]
)
onnx.save(model, 'model_modified.onnx')
注意：手动修改ONNX模型风险很高。改错一个维度，整个计算图可能就废了。我建议只在万不得已时用这招，而且改完后一定要做数值验证。
4.5 算子兼容性检查
模型转成ONNX只是第一步。你最终要把它部署到某个推理引擎上——比如TensorRT、NCNN、OpenVINO。这些引擎支持的ONNX算子集各不相同。如果模型里用了某个引擎不支持的算子，转换就会失败。
4.5.1 常见不兼容算子
| 算子 | TensorRT支持 | NCNN支持 | OpenVINO支持 |
|---|---|---|---|
| Resize (nearest) | ✓ | ✓ | ✓ |
| Resize (linear) | ✓ (需特定版本) | ✓ | ✓ |
| NonMaxSuppression | ✓ (TRT 8+) | ✗ | ✓ |
| ScatterND | ✗ | ✗ | ✓ |
| Einsum | ✗ | ✗ | ✗ |
4.5.2 如何检查兼容性
我一般用三步走：
- 用onnxruntime跑一遍——确认ONNX本身没问题。
- 用目标引擎的ONNX解析器跑一遍——比如TensorRT的
trtexec，看它能不能解析。 - 逐算子检查——如果解析失败，用
onnx库遍历所有节点，对照目标引擎的算子支持列表，找出不兼容的算子。
import onnx
model = onnx.load('model.onnx')
unsupported_ops = []
# 假设这是TensorRT不支持的算子列表
trt_unsupported = ['Einsum', 'ScatterND', 'NonMaxSuppression']
for node in model.graph.node:
if node.op_type in trt_unsupported:
unsupported_ops.append(node.op_type)
if unsupported_ops:
print(f"不兼容算子: {set(unsupported_ops)}")
else:
print("所有算子均兼容")
核心思路：算子兼容性检查不是一次性的。你换一个推理引擎版本，支持列表可能就变了。我建议把检查脚本写进CI流程，每次模型更新自动跑一遍。
4.6 本章知识结构图
下面这张图概括了模型转换到ONNX再到推理引擎的完整流程，以及每个环节的关键检查点。
4.7 本章小结
模型转换这件事，说难不难，说简单也不简单。核心就三点：
- 导出参数要设对——opset版本、dynamic_axes、输入输出名称，一个都不能错。
- 结构要能看懂——学会用netron和onnx库解析模型，出了问题能定位。
- 算子要提前查——别等到部署时才发现某个算子不支持，那时候改模型就晚了。
我做了这么多项目，最深的体会是：模型转换花的时间，往往比训练还多。但只要你把这一章的内容吃透了，后面部署到TensorRT、NCNN这些引擎上，就会顺畅很多。