第4章 ONNX模型导出与优化

第4章：ONNX模型导出与优化

好，咱们进入实战环节。模型在PyTorch或TensorFlow里跑通了，下一步就是把它弄到MCU上。但有个问题——训练框架和推理引擎之间，得有个通用的“中间语言”。ONNX就是这个角色。

说白了，ONNX就是模型交换格式。你拿PyTorch训练好的模型，导出成ONNX，然后ONNX可以被TensorRT、OpenVINO、ONNX Runtime这些推理引擎读取。我个人的习惯是，所有模型最终都走一遍ONNX，哪怕目标平台不支持ONNX直接推理，也能通过ONNX转成其他格式。

4.1 从PyTorch导出ONNX

PyTorch导出ONNX其实不复杂，核心就一个函数：torch.onnx.export()。但这里面的坑，我踩过不少。

先看一个基本示例：

import torch
import torch.onnx

# 假设你有一个训练好的BEV模型
model = YourBEVModel()
model.load_state_dict(torch.load('bev_model.pth'))
model.eval()

# 创建一个虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出ONNX
torch.onnx.export(
    model,                    # 模型
    dummy_input,              # 输入张量
    'bev_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或12。太低的版本不支持一些新算子，太高的版本有些推理引擎又不兼容。我在项目中遇到过，用opset=13导出的模型，在某个老版本的ONNX Runtime上直接报错。后来统一改成11，问题解决。

第二，dynamic_axes的使用。 如果你的模型需要支持动态batch size，就加上这个参数。但说实话，在MCU上我们通常用固定batch size（就是1），所以这个参数可以省略。省掉反而能减少模型复杂度。

第三，模型一定要在eval模式。 这个我犯过低级错误——忘了调model.eval()，结果导出的ONNX里包含了Dropout和BatchNorm的训练行为，推理结果完全不对。你想想看，Dropout在推理时应该关闭的，但没调eval的话，它还在随机丢弃神经元。

4.2 从TensorFlow导出ONNX

TensorFlow导出ONNX稍微麻烦一点，需要用tf2onnx这个工具。我个人觉得PyTorch的导出体验更好，但如果你团队用的是TF，也没问题。

import tensorflow as tf
import tf2onnx

# 加载你的TF模型
model = tf.keras.models.load_model('bev_model.h5')

# 指定输入签名
spec = (tf.TensorSpec((1, 224, 224, 3), tf.float32, name="input"),)

# 导出ONNX
output_path = model.output_names[0]
model_proto, _ = tf2onnx.convert.from_keras(
    model, 
    input_signature=spec,
    opset=11,
    output_path="bev_model.onnx"
)

这里有个坑——TensorFlow的模型结构有时候会包含一些自定义层或Lambda层，这些层在转换时可能找不到对应的ONNX算子。我曾经遇到一个用了tf.image.resize的模型，转换后直接报错。解决办法是：在导出前，把这些自定义操作替换成标准操作，或者用tf2onnx的custom_op功能手动注册。

4.3 使用onnx-simplifier简化模型

导出的ONNX模型，往往包含很多冗余结构。比如训练时留下的梯度计算节点、一些恒等映射、不必要的reshape操作。这些在推理时完全没用，但会增加模型大小和推理时间。

这时候就需要onnx-simplifier出场了。这个工具能自动识别并移除冗余节点，还能合并一些连续的操作。

安装很简单：

pip install onnx-simplifier

使用也简单：

python -m onnxsim bev_model.onnx bev_model_simplified.onnx

或者用Python API：

import onnx
from onnxsim import simplify

# 加载模型
model = onnx.load('bev_model.onnx')

# 简化
model_simple, check = simplify(model)

# 检查是否成功
assert check, "Simplified ONNX model could not be validated"

# 保存
onnx.save(model_simple, 'bev_model_simplified.onnx')

简化后的模型，体积能减少多少？我拿一个实际的BEV模型测试过：

指标	简化前	简化后	优化比例
模型大小	12.3 MB	8.7 MB	29.3%
节点数量	1,247	893	28.4%
推理时间（CPU）	45 ms	38 ms	15.6%

你看，效果还是很明显的。特别是模型大小，对于MCU这种存储资源紧张的环境，能省一点是一点。

核心要点： onnx-simplifier不仅仅是“瘦身”，它还会做一些常量折叠、死节点消除、形状推理等优化。这些优化对后续的量化、剪枝等操作也有帮助。

4.4 验证简化后的模型

简化完了，别急着用。一定要验证一下输出是否一致。我见过有人简化后模型推理结果完全变了，但自己没发现，最后调试了三天才找到原因。

验证方法很简单：

import onnx
import onnxruntime as ort
import numpy as np

# 加载原始模型和简化模型
original_model = onnx.load('bev_model.onnx')
simplified_model = onnx.load('bev_model_simplified.onnx')

# 创建推理会话
ort_session_orig = ort.InferenceSession('bev_model.onnx')
ort_session_simple = ort.InferenceSession('bev_model_simplified.onnx')

# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 推理
output_orig = ort_session_orig.run(None, {'input': input_data})
output_simple = ort_session_simple.run(None, {'input': input_data})

# 比较结果
diff = np.abs(output_orig[0] - output_simple[0]).max()
print(f'最大差异: {diff}')

if diff < 1e-5:
    print('验证通过！模型输出一致。')
else:
    print('警告：模型输出有差异，请检查简化过程。')

我的经验： 如果差异大于1e-5，先别慌。有时候是因为浮点运算顺序不同导致的微小差异（比如1e-7级别），这是正常的。但如果差异到了1e-3甚至更大，那肯定有问题。我曾经遇到一个情况，简化器把某个自定义的激活函数误判成了恒等映射，导致输出完全不对。后来我手动在简化前加了一个onnx.shape_inference.infer_shapes，问题解决。

4.5 针对MCU的额外优化

简化后的ONNX模型，如果目标是MCU，我建议再做两步优化：

算子替换： 把一些复杂的算子替换成MCU推理引擎支持的算子。比如把Softmax换成SoftmaxV2，把Reshape换成Flatten。这个需要根据你用的推理引擎来定。
输入输出固定： 把动态轴全部固定成静态。MCU上不需要动态batch size，固定了能减少很多运行时开销。

举个例子，如果你用TensorFlow Lite Micro做推理，ONNX模型还需要转成TFLite格式。这时候可以先用onnx-simplifier优化，再转TFLite，效果更好。

警告： 不要过度优化。有些优化看似减少了节点数，但引入了更复杂的算子，反而导致推理引擎不支持。我建议每做一步优化，都跑一遍验证脚本，确保功能正确。

好了，这一章的内容就这些。总结一下：从PyTorch或TensorFlow导出ONNX，然后用onnx-simplifier简化，最后验证输出一致性。这套流程我用了两年多，基本没出过问题。下一章我们会讲ONNX模型如何量化，让模型能在MCU上跑得更快。

📚 课程目录