4、ONNX模型优化：ONNX Simplifier使用、常量折叠、节点融合、精度选择（FP32/FP16/INT8）

好，咱们进入第四章。ONNX模型优化，这步其实挺关键的。

你想想看，从PyTorch或者TensorFlow导出的ONNX模型，很多时候都带着一堆冗余操作。就像你搬家，把所有东西都塞进箱子，但到了新家，总得拆开整理一下，扔掉没用的，把零碎的东西组合起来。ONNX优化，干的就是这个活。

4.1 为什么要优化ONNX模型？

我刚开始做部署的时候，也犯过懒。觉得模型能跑起来不就行了？结果有一次，一个ResNet50模型，导出后直接扔给TensorRT，发现推理速度比预期慢了不少。后来一查，发现ONNX图里有一堆没用的Cast节点和Shape操作。嗯，这就是没优化的后果。

优化ONNX模型，主要解决三个问题：

  • 减少冗余计算：去掉那些训练时有用、推理时没用的节点，比如Dropout、Identity操作。
  • 简化图结构：把常量折叠掉，把连续的线性操作融合成一个节点。
  • 适配目标硬件：为后续的TensorRT转换铺平道路，避免不支持的操作导致转换失败。

核心原则：优化后的ONNX模型，在数学上必须和原始模型完全等价。不能因为优化而改变了模型的输出结果。

4.2 ONNX Simplifier：一键简化神器

ONNX Simplifier，简称onnxsim，是我用得最多的工具。说白了，它就是个自动化的模型清理工。

安装很简单：

pip install onnx-simplifier

使用更简单，一行命令：

python -m onnxsim input_model.onnx output_model.onnx

它会帮你做哪些事呢？

  • 常量折叠：把计算图中所有能用常量算出来的部分，提前算好。比如一个Add操作，两个输入都是常量，那直接算出一个新常量替换掉这个节点。
  • 节点消除：去掉那些对结果没影响的节点，比如Identity、Dropout（推理模式下）。
  • 形状推理：把动态的形状信息尽量变成静态的，方便后续优化。

我的经验：我个人习惯在导出ONNX后，第一时间就跑一遍onnxsim。它解决了我遇到过的大部分“模型转TensorRT报错”的问题。尤其是那些因为动态Shape导致的错误，onnxsim常常能帮你搞定。

4.3 常量折叠：把能算的都提前算好

常量折叠，听起来高大上，其实原理很简单。就是编译器里的一个经典优化技术。

举个例子，你的模型里有一段这样的操作：

# 伪代码
weight = np.array([1.0, 2.0, 3.0])
bias = np.array([0.1, 0.2, 0.3])
output = weight + bias  # 这个结果在推理时永远不会变

那为什么不直接在导出模型时，就把output算出来，存成一个常量呢？

这样做的好处很明显：

  • 减少计算量：推理时少执行一个Add操作。
  • 减少内存访问：不需要反复读取weight和bias。
  • 简化图结构：让后续的节点融合更容易。

我曾经遇到过一个模型，里面有一大堆Gemm（矩阵乘法）操作，但权重矩阵全是固定的。用onnxsim跑完后，模型体积直接缩小了30%，推理速度也快了将近10%。

注意：常量折叠不是万能的。如果常量太大（比如一个巨大的权重矩阵），折叠后反而会让模型文件变大。不过这种情况比较少见，大多数时候利大于弊。

4.4 节点融合：把零碎操作打包

节点融合，就是把多个连续的小操作，合并成一个大的操作。这有点像把散落在地上的零件组装成一个模块。

常见的融合模式有：

  • Conv + BN + ReLU：这是最经典的融合模式。卷积、批归一化、ReLU激活，三个操作可以融合成一个Conv节点。TensorRT内部会自动做这个优化，但我们在ONNX层面提前做好，可以让转换更顺畅。
  • Add + ReLU：残差网络里常见的模式。
  • MatMul + Add：全连接层加偏置。

为什么需要节点融合？

  • 减少Kernel Launch开销：每个操作都需要启动一个CUDA Kernel。融合后，启动次数减少，延迟降低。
  • 减少显存带宽占用：中间结果不需要写回显存，直接在寄存器里传递。

避坑指南：我曾经在融合Conv + BN时，因为BN层的参数精度问题，导致融合后的模型精度下降了0.1%。后来发现是FP16精度下，BN的均值和方差被截断了。解决办法是：在融合前，先把BN层的参数转成FP32，融合完再转回去。

4.5 精度选择：FP32、FP16、INT8

精度选择，说白了就是在速度和精度之间做权衡。你想想看，数值越精确，需要的位数就越多，计算就越慢。反过来，位数越少，计算越快，但精度可能会下降。

三种精度的对比：

精度类型 位数 相对速度 精度损失 适用场景
FP32 32位 1x（基准） 无 精度要求极高的场景，如医疗影像
FP16 16位 约2x 通常可忽略 大多数推理场景，兼顾速度和精度
INT8 8位 约4x 需要校准，可能明显 对延迟极度敏感的场景，如实时视频分析

4.5.1 FP32：最稳妥的选择

FP32是单精度浮点数。如果你对精度有100%的信心，或者模型很小，推理速度不是瓶颈，那就用FP32。我刚开始做部署时，所有模型都用FP32，因为不想引入任何不确定性。

4.5.2 FP16：性价比之王

FP16是半精度浮点数。现在NVIDIA的Turing架构及以后的GPU，都对FP16有很好的硬件支持。速度能翻倍，精度损失通常可以忽略不计。

怎么在ONNX里设置FP16？其实ONNX本身不直接支持FP16的权重存储。我们通常是在转TensorRT时，指定使用FP16精度。但有一点要注意：如果你的ONNX模型里有一些操作对精度特别敏感（比如某些LayerNorm），FP16可能会导致数值不稳定。

我的建议：对于大多数视觉模型（分类、检测、分割），FP16是首选。我做过一个YOLOv5的部署，FP16比FP32快了将近一倍，而mAP只下降了0.2%。这个代价完全可以接受。

4.5.3 INT8：极限压榨性能

INT8是8位整型。速度最快，但精度损失也最大。而且，INT8需要做校准（Calibration）。

校准的过程是这样的：

  1. 准备一批有代表性的校准数据（通常几百到几千张图片）。
  2. 用FP32模型跑一遍，收集每一层的激活值分布。
  3. 根据分布，计算出从FP32到INT8的缩放因子。
  4. 用这个缩放因子，把模型量化成INT8。

校准数据的选择很关键。如果校准数据跟实际推理数据分布不一致，量化后的模型精度可能会崩掉。

我曾经踩过的坑：有一次做INT8量化，用了ImageNet的验证集做校准，结果模型在真实业务数据上精度掉了5%。后来换成业务数据做校准，精度才恢复正常。所以，校准数据一定要用跟实际场景最接近的数据。

4.6 优化流程总结

好了，咱们把这一章的内容串起来。一个典型的ONNX优化流程是这样的：

  1. 导出ONNX：从训练框架导出原始ONNX模型。
  2. 运行onnxsim：一键简化，去掉冗余，折叠常量。
  3. 手动检查：用Netron可视化工具，看看图里还有没有奇怪的节点。如果有，可能需要手动写脚本处理。
  4. 精度选择：根据业务需求，决定用FP32、FP16还是INT8。
  5. 转TensorRT：把优化后的ONNX模型，转成TensorRT引擎。

记住，优化不是一步到位的。你可能需要反复调整，才能找到最适合你模型和硬件的方案。嗯，这就是工程实践的魅力所在。

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