4. ONNX模型优化：常量折叠、节点融合、图优化、使用onnx-simplifier简化模型
模型从训练框架导出成ONNX，这步其实不难。真正让我头疼的，是导出来的ONNX模型又大又慢，跑在端侧根本没法用。我刚开始做部署时，就吃过这个亏——模型在PC上推理好好的，一放到开发板上，延迟直接翻了三倍。
后来我才明白，训练框架导出的ONNX，里面藏着大量冗余计算。说白了，就是训练框架为了灵活性，加了很多“中间人”节点。这些节点在训练时有用，推理时就是累赘。所以，ONNX模型优化，就是把这些累赘去掉。
这一章，我带你走一遍ONNX优化的核心流程。包括常量折叠、节点融合、图优化，还有我最常用的工具——onnx-simplifier。
4.1 为什么要做ONNX优化？
先看一个真实案例。我之前优化过一个BEV感知模型，原始ONNX有1200多个节点。用onnx-simplifier跑一遍，直接砍到400多个节点。推理速度提升了40%。
你想想看，端侧设备的算力本来就有限。多一个节点，就多一次内存读写，多一次调度开销。优化ONNX，本质上就是给模型“减负”。
核心目标：
- 减少节点数——降低推理引擎的调度开销
- 消除冗余计算——减少不必要的算子执行
- 合并连续操作——减少内存读写次数
- 简化图结构——提高后续量化/编译的稳定性
4.2 常量折叠（Constant Folding）
常量折叠，说白了就是把那些输入全是常量的计算，提前算好。比如一个Add节点，两个输入都是固定数值，那推理时每次都要算一遍，纯属浪费。
我举个例子。模型里有个Shape节点，后面跟着Gather取某个维度。如果输入张量的形状在导出时就已经固定了，那这个Shape+Gather的结果就是个常量。常量折叠会把这两个节点替换成一个Constant节点。
# 常量折叠前后的对比
# 优化前：Shape -> Gather -> 后续节点
# 优化后：Constant -> 后续节点ONNX Runtime自带的优化器就支持常量折叠。但要注意，不是所有常量都能被折叠。如果某个常量计算依赖动态输入，那就没法折叠。嗯，这里要区分清楚。
我的经验：常量折叠对BEV模型的GridSample和Resize操作特别有效。因为这些操作经常用到固定的坐标网格，折叠后能省掉一大半计算量。
4.3 节点融合（Node Fusion）
节点融合，就是把多个连续的算子合并成一个。这样做的好处很明显：减少内存读写，减少调度开销。
最常见的融合模式是Conv + BatchNorm + Relu。训练时这三个是分开的，推理时完全可以合并成一个Conv节点。BN的参数可以吸收到Conv的权重和偏置里。
| 融合模式 | 原始节点序列 | 融合后节点 | 收益 |
|---|---|---|---|
| Conv-BN-ReLU | Conv → BatchNorm → Relu | Conv | 减少2次内存读写 |
| MatMul-Add | MatMul → Add | Gemm | 减少1次内存读写 |
| Reshape-Transpose-Reshape | Reshape → Transpose → Reshape | Transpose（优化后） | 减少2次reshape开销 |
| Clip-Sigmoid | Clip → Sigmoid | Sigmoid（输入已裁剪） | 减少1次clip计算 |
我遇到过最夸张的一次，一个BEV模型的backbone部分，通过Conv-BN-ReLU融合，节点数从600降到了200。推理速度直接翻倍。
注意：节点融合不是万能的。有些融合后的算子，端侧推理引擎可能不支持。比如某些NPU不支持融合后的Conv（带BN参数）。这时候就得权衡：是保留原始结构，还是换一种融合方式。
4.4 图优化（Graph Optimization）
图优化，范围比常量折叠和节点融合更广。它是对整个计算图做结构性调整。常见的图优化包括：
- 死节点消除——删除那些输出没有被使用的节点
- 公共子表达式消除——如果两个节点计算相同，合并成一个
- 恒等操作消除——比如
Identity节点，直接删掉 - Cast消除——如果数据类型转换是多余的，直接去掉
- Shape推理优化——提前推理出中间张量的形状，减少动态shape判断
我个人习惯用ONNX Runtime的GraphOptimizationLevel来控制优化强度。有三个级别：
# ONNX Runtime的图优化级别
# ORT_ENABLE_BASIC (1) - 基础优化，常量折叠 + 死节点消除
# ORT_ENABLE_EXTENDED (2) - 扩展优化，包含节点融合
# ORT_ENABLE_ALL (99) - 全部优化，包含所有已知优化我建议在端侧部署时，先用ORT_ENABLE_ALL跑一遍。如果推理引擎报错说某个算子不支持，再降级到ORT_ENABLE_EXTENDED。
4.5 使用onnx-simplifier简化模型
onnx-simplifier是我最常用的工具，没有之一。它把常量折叠、节点融合、图优化打包在一起，一行命令就能搞定。
它的核心思路是：用ONNX Runtime跑一次推理，记录下所有中间结果。然后对比原始图和推理结果，把那些“计算前后结果不变”的节点删掉。
# 安装onnx-simplifier
pip install onnx-simplifier
# 简化模型
python -m onnxsim input_model.onnx output_model.onnx
# 带输入shape的简化（推荐）
python -m onnxsim input_model.onnx output_model.onnx \
--input-shape "input:1,3,224,224"我曾经用onnx-simplifier优化过一个BEV模型的ViewTransformer部分。原始模型里有大量Reshape和Transpose操作，简化后直接减少了30%的节点数。
使用技巧：
- 一定要指定
--input-shape，否则简化器不知道输入的具体形状，很多优化做不了 - 简化后，用
onnx.checker.check_model验证一下模型是否合法 - 如果简化后的模型推理结果不对，试试加
--dynamic-input-shape参数
4.6 优化流程实战
下面是我在实际项目中总结的ONNX优化流程。你照着做，基本不会踩坑。
- 导出ONNX——从PyTorch/TensorFlow导出原始ONNX模型
- 运行onnx-simplifier——指定输入shape，做第一轮简化
- ONNX Runtime优化——用
ORT_ENABLE_ALL做第二轮优化 - 验证精度——对比优化前后的推理结果，误差应在1e-5以内
- 端侧测试——在目标设备上跑一遍，确认算子都支持
- 迭代调整——如果某个算子不支持，回退到上一版，手动调整
我曾经在一个项目中，因为跳过了第4步，直接拿优化后的模型去部署。结果在端侧推理结果全是NaN。排查了半天，才发现是某个融合后的算子精度不达标。从那以后，我再也不敢跳过精度验证了。
4.7 优化效果对比
拿一个典型的BEV感知模型（基于ResNet-50的backbone + Transformer的view transformer）做测试：
| 优化阶段 | 节点数 | 模型大小 | 推理延迟（RK3588） |
|---|---|---|---|
| 原始ONNX | 1247 | 187 MB | 142 ms |
| onnx-simplifier | 683 | 152 MB | 98 ms |
| + ORT优化 | 412 | 148 MB | 85 ms |
| + 手动融合 | 387 | 146 MB | 79 ms |
你看，从1247个节点降到387个，延迟从142ms降到79ms。优化效果非常明显。
4.8 ONNX优化流程图
下面这张图，展示了ONNX优化的完整流程。我建议你把它打印出来，贴在工位上。
4.9 常见问题与避坑
最后，我总结几个常见问题。这些都是我踩过的坑，你遇到了可以直接拿来用。
问题1：简化后的模型推理结果不对
原因：onnx-simplifier在简化时，如果输入shape是动态的，可能会误删某些节点。
解决方案：加--dynamic-input-shape参数，或者手动指定所有输入的具体shape。
问题2：端侧推理引擎报错“Unsupported Operator”
原因：节点融合后产生了端侧不支持的算子。
解决方案：回退到未融合的版本，手动实现一个等效的算子序列。
问题3：优化后模型变小了，但推理反而变慢了
原因：某些融合后的算子，在端侧没有硬件加速，只能用CPU跑。
解决方案：对比不同融合策略的性能，选择最优方案。
我曾经在一个项目中，优化后的模型在PC上推理快了30%，但到了RK3588上反而慢了20%。排查后发现，融合后的Conv算子（带BN参数）在NPU上不支持，只能回退到CPU。最后我手动把BN参数拆出来，用CPU跑BN，NPU跑Conv，性能才恢复正常。
嗯，ONNX优化这块，说白了就是“实践出真知”。多试几次，你就能找到最适合自己模型的优化方案。