4、ONNX模型导出与转换：从PyTorch/TensorFlow导出ONNX模型、ONNX模型结构分析、模型验证工具。
好，咱们进入实战环节。前面聊了那么多ONNX Runtime的底层原理，现在终于要动手了。这一章，我带你走一遍模型导出的完整流程。说白了，就是把你在PyTorch或TensorFlow里训好的模型，变成ONNX格式。这一步要是没走对，后面所有的优化都是白搭。
4.1 从PyTorch导出ONNX模型
PyTorch导出ONNX，我个人觉得是最顺手的。它原生支持`torch.onnx.export`接口，几乎不需要额外装什么库。但这里有个坑——动态轴的处理。
先看一个最简单的例子。假设你有一个分类模型，输入是固定尺寸的图片：
import torch
import torchvision.models as models
# 加载一个预训练的ResNet18
model = models.resnet18(pretrained=True)
model.eval()
# 创建一个虚拟输入，batch_size=1，3通道，224x224
dummy_input = torch.randn(1, 3, 224, 224)
# 导出ONNX
torch.onnx.export(
model, # 模型
dummy_input, # 输入张量
"resnet18.onnx", # 输出文件名
export_params=True, # 导出训练好的参数
opset_version=11, # ONNX算子集版本
do_constant_folding=True, # 常量折叠优化
input_names=['input'], # 输入名称
output_names=['output'], # 输出名称
dynamic_axes={
'input': {0: 'batch_size'}, # 动态batch
'output': {0: 'batch_size'}
}
)
print("导出成功！")
嗯，这里要注意`dynamic_axes`参数。我在项目中遇到过，很多同学忘了设置这个，结果导出的模型只能跑固定batch size。你想想看，部署到ARM设备上，有时候一次只推理一张图，有时候想批量处理，没有动态轴就尴尬了。
4.2 从TensorFlow导出ONNX模型
TensorFlow导出ONNX稍微麻烦一点。你需要用`tf2onnx`这个工具。我个人习惯用SavedModel格式作为中间桥梁。
步骤其实不复杂：
- 先把TensorFlow模型保存为SavedModel格式
- 然后用tf2onnx命令行工具转换
看代码：
# 第一步：保存TensorFlow模型为SavedModel
import tensorflow as tf
# 假设你有一个训练好的模型
model = tf.keras.applications.MobileNetV2(weights='imagenet')
model.save('mobilenetv2_savedmodel')
# 第二步：用tf2onnx转换
# 命令行执行：
# python -m tf2onnx.convert \
# --saved-model mobilenetv2_savedmodel \
# --output mobilenetv2.onnx \
# --opset 11
这里有个小技巧。我建议你在转换前，先用`onnxruntime`的Python接口跑一下原始模型，确认输入输出的shape和dtype。为什么？因为TensorFlow的图有时候会包含一些奇怪的算子，转换时容易出问题。
4.3 ONNX模型结构分析
模型导出来之后，别急着部署。先看看它长什么样。我常用的工具是`Netron`和`onnxruntime`自带的模型分析接口。
Netron是个可视化工具，直接拖拽ONNX文件进去就能看到网络结构。但说实话，在ARM设备上开发时，我更习惯用命令行分析：
import onnx
# 加载模型
model = onnx.load("resnet18.onnx")
# 查看模型基本信息
print("IR版本:", model.ir_version)
print("生产者:", model.producer_name)
print("图名称:", model.graph.name)
# 查看输入输出
for input in model.graph.input:
print(f"输入: {input.name}, shape: {input.type.tensor_type.shape}")
for output in model.graph.output:
print(f"输出: {output.name}, shape: {output.type.tensor_type.shape}")
# 统计节点数量
print(f"节点总数: {len(model.graph.node)}")
# 查看算子类型分布
op_types = {}
for node in model.graph.node:
op_type = node.op_type
op_types[op_type] = op_types.get(op_type, 0) + 1
print("算子分布:", op_types)
你想想看，这一步能帮你发现很多问题。比如，我曾在导出的模型里发现了一个`Identity`节点，它什么都不做，但白白增加了推理延迟。这种冗余节点，在ARM设备上就是性能杀手。
| 分析项 | 说明 | 常见问题 |
|---|---|---|
| IR版本 | ONNX中间表示版本 | 版本过低可能不支持新算子 |
| 输入输出 | 检查shape和dtype | 动态轴未设置、数据类型不匹配 |
| 算子分布 | 统计各类算子数量 | 出现不支持的自定义算子 |
| 节点总数 | 网络中的操作节点数 | 冗余节点过多影响性能 |
4.4 模型验证工具
模型结构没问题了，接下来要验证精度。说白了，就是确保ONNX模型的推理结果和原始框架一致。
我常用的验证流程分三步：
- 数值一致性检查：用同样的输入，对比ONNX Runtime和原始框架的输出
- 精度损失评估：在验证集上跑一遍，看精度下降了多少
- 性能基准测试：在目标ARM设备上测推理延迟
先看数值一致性检查的代码：
import onnxruntime as ort
import numpy as np
import torch
# 原始PyTorch模型推理
model.eval()
with torch.no_grad():
torch_output = model(torch.from_numpy(dummy_input)).numpy()
# ONNX Runtime推理
ort_session = ort.InferenceSession("resnet18.onnx")
ort_inputs = {ort_session.get_inputs()[0].name: dummy_input}
ort_output = ort_session.run(None, ort_inputs)[0]
# 对比结果
diff = np.abs(torch_output - ort_output)
print(f"最大绝对误差: {diff.max():.6f}")
print(f"平均绝对误差: {diff.mean():.6f}")
# 检查是否在容忍范围内
if diff.max() < 1e-3:
print("✅ 数值一致性通过！")
else:
print("❌ 数值差异过大，请检查模型导出过程")
精度损失评估这一步，我建议用你训练时的验证脚本，把模型替换成ONNX版本跑一遍。我在项目中遇到过，ONNX模型在ImageNet上top-1精度只掉了0.1%，但换到自己的业务数据集上，直接掉了2%。原因是什么？数据分布不同，某些算子对特定数据更敏感。
最后是性能基准测试。在ARM设备上，我一般测100次取平均：
import time
# 预热
for _ in range(10):
ort_session.run(None, ort_inputs)
# 正式测试
times = []
for _ in range(100):
start = time.time()
ort_session.run(None, ort_inputs)
times.append(time.time() - start)
print(f"平均推理时间: {np.mean(times)*1000:.2f} ms")
print(f"标准差: {np.std(times)*1000:.2f} ms")
嗯，这里要注意。ARM设备上CPU频率可能会动态调整，导致测试结果不稳定。我习惯先锁定CPU频率，或者跑足够多次取中位数。另外，内存带宽也是瓶颈，多线程推理时尤其明显。
好了，模型导出、结构分析、验证这三步走完，你的ONNX模型就可以放心部署到ARM设备上了。下一章，我会讲怎么用ONNX Runtime的Session Options做性能调优，到时候咱们再细聊。