4、ONNX模型优化：ONNX Simplifier使用、常量折叠、节点融合、冗余消除、ONNX opset版本选择

好，咱们接着聊ONNX模型优化。说实话，很多同学把模型从PyTorch或者TensorFlow转到ONNX之后，就觉得万事大吉了。其实不然。你导出的ONNX模型，就像刚写完的代码——能跑，但跑得不一定快，甚至可能藏着一些“垃圾代码”。

我刚开始做TensorRT部署那会儿，就吃过这个亏。一个看起来挺简单的分类模型，转成TensorRT之后死活跑不到预期帧率。后来一查，ONNX图里有一堆没用的Identity节点，还有几个可以合并的卷积+BN没合并。嗯，从那以后，我养成了一个习惯：ONNX导出之后，第一件事就是做优化。

4.1 为什么要优化ONNX模型？

说白了，框架导出的ONNX模型，往往不是最优的。它可能包含很多冗余操作，比如：

  • 常量折叠没做：一些在推理时不会变化的计算，比如Shape、Gather、Unsqueeze的组合，完全可以提前算好。
  • 节点太碎：一个简单的卷积+批归一化，在ONNX里可能被拆成好几个小节点。TensorRT虽然能处理，但效率会打折扣。
  • 有死节点：有些节点计算出来的结果，后面根本没人用。这些就是“垃圾”，得清理掉。
  • opset版本太老或太新：版本不对，有些算子TensorRT不支持，或者支持得不好。

你想想看，带着这些“包袱”去构建TensorRT引擎，能快才怪。所以，优化ONNX模型，是搭建高性能推理引擎的必修课。

4.2 ONNX Simplifier：一键优化的利器

ONNX Simplifier，简称onnxsim，是我最常用的工具之一。它就像一个“代码格式化+代码压缩”工具，能帮你自动完成大部分常规优化。

安装很简单：

pip install onnx-simplifier

使用更简单，一行命令：

python -m onnxsim input_model.onnx output_model.onnx

它会做什么呢？我拆开来说：

  • 常量折叠（Constant Folding）：把那些输入全是常量的子图，直接计算成结果。比如一个Shape操作，输入张量的形状在模型固定后就是已知的，没必要每次推理都算一遍。
  • 节点融合（Node Fusion）：把一些可以合并的节点合并成一个。比如Conv + BatchNormalization，可以融合成一个Conv节点（权重和偏置重新计算）。
  • 冗余消除（Redundancy Elimination）：删除那些输出没有被任何节点使用的“死节点”，以及删除一些恒等映射（Identity）节点。

我个人习惯，每次导出ONNX后，第一件事就是跑一遍onnxsim。它解决了我80%的优化问题。

小提示： onnxsim默认会做“形状推断”（shape inference），这能帮助它更好地进行常量折叠。如果你的模型有动态形状，记得加上--dynamic-input-shape参数，否则可能会把动态维度也给折叠了，导致模型失效。

4.3 深入理解：常量折叠、节点融合与冗余消除

虽然onnxsim一键搞定，但我还是建议你理解一下它背后的原理。这样遇到复杂情况时，你才知道怎么手动处理。

4.3.1 常量折叠

举个例子，你的模型里可能有这么一段：

// 伪代码
shape = Shape(input_tensor)  // 获取形状 [1, 3, 224, 224]
gather = Gather(shape, indices=[2]) // 取出第2个维度，得到 [224]
unsqueeze = Unsqueeze(gather, axes=[0]) // 变成 [1, 224]

如果input_tensor的形状在模型编译时就是确定的，那这段代码完全可以用一个常量[1, 224]替代。这就是常量折叠。

我在项目中遇到过，一个Transformer模型里，这种“取形状-切片-拼接”的操作反复出现。onnxsim跑完之后，模型体积直接小了30%，推理速度也快了5%。

4.3.2 节点融合

最经典的融合就是Conv + BN。BN在推理时，其实就是一个线性变换：y = gamma * (x - mean) / sqrt(var + eps) + beta。而卷积也是一个线性变换。两个线性变换叠加，完全可以合并成一个新的卷积。

合并后的卷积，权重和偏置会重新计算：

// 新权重 = 旧权重 * (gamma / sqrt(var + eps))
// 新偏置 = (旧偏置 - mean) * (gamma / sqrt(var + eps)) + beta

这样做的好处很明显：少了一次内存访问，少了一次kernel launch，速度自然就上去了。

核心要点： 节点融合的本质，是把多个计算密集型的操作合并成一个，减少数据在GPU显存和计算单元之间的搬运次数。

4.3.3 冗余消除

这个最好理解。有些节点，比如Identity（恒等映射），它的输出等于输入，纯粹是为了方便框架内部做梯度计算或者调试用的。在推理时，这些节点完全可以删掉，直接把输入接到输出上。

还有一些节点，它的输出没有被任何其他节点引用。这种“死节点”也要删掉。我曾经在一个导出的ONNX模型里，发现了一整段“死代码”——一个分支计算的结果，在后面的图中根本没用到。估计是导出时某个条件判断留下的残骸。

4.4 ONNX Opset版本选择：一个容易踩的坑

ONNX的opset，你可以理解为算子的“版本号”。每个新版本都会增加一些新算子，或者修改旧算子的行为。

为什么opset版本很重要？因为TensorRT对ONNX算子的支持，是跟opset版本挂钩的。比如：

  • TensorRT 8.x 一般推荐使用 opset 13 或 16。
  • TensorRT 10.x 可能支持到 opset 20 甚至更高。

选错了opset会怎样？

  • 版本太低：有些新算子（比如Einsum、LayerNormalization）可能没有，或者只能用一堆旧算子拼凑，效率低下。
  • 版本太高：TensorRT可能还没实现这个版本的某些算子，导致转换失败。

我建议的做法是：

  1. 先查一下你用的TensorRT版本，官方文档里会写明它支持的最高opset版本。
  2. 导出ONNX时，把opset设置成这个最高版本，或者比它低1-2个版本（更稳妥）。
  3. 如果遇到不支持的算子，再考虑降级opset，或者用onnx-simplifier的--skip-optimization选项，手动处理有问题的部分。

在PyTorch中导出时，可以这样指定：

import torch
# 假设你的模型是 model
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    opset_version=16,  # 这里指定opset版本
    input_names=["input"],
    output_names=["output"]
)
避坑指南： 我曾经在TensorRT 8.4上，用opset 17导出了一个包含ScatterND算子的模型。结果TensorRT直接报错，说这个算子不支持。后来我把opset降到13，用ScatterElements替代，才顺利跑通。所以，不要盲目追求最新opset，兼容性才是第一位的。

4.5 优化流程总结

好了，说了这么多，我总结一下我个人的ONNX优化流程：

  1. 导出：用框架导出ONNX，opset版本选TensorRT支持的最高版本或低1-2个版本。
  2. 简化：用onnxsim跑一遍，完成常量折叠、节点融合、冗余消除。
  3. 检查：用Netron（一个ONNX可视化工具）打开优化后的模型，看看图结构是否干净。如果还有明显的冗余，手动处理。
  4. 验证：用onnxruntime跑一下优化前后的模型，确保精度一致。

这一步做完，你的ONNX模型就“干净”多了。接下来再交给TensorRT，它构建引擎的速度会更快，生成的引擎性能也会更好。

嗯，关于ONNX模型优化，就先聊这么多。下一节，我们进入TensorRT的核心——网络定义与引擎构建。到时候你会发现，一个干净的ONNX模型，能让你省去很多麻烦。

← 上一章 📖 返回目录 下一章 →