3. 算子融合与图优化

图优化，说白了就是让模型跑得更快、更省资源。你想想看，我们辛辛苦苦训练出来的模型，如果直接扔到手机或者嵌入式设备上，那效率简直惨不忍睹。我刚开始做端侧推理那会儿，就吃过这个亏——模型在服务器上跑得飞快，一上手机，直接卡成PPT。

后来我才明白，图优化是端侧部署的必修课。今天我们就来聊聊三个最核心的优化手段：算子融合、常量折叠和死节点消除。

3.1 算子融合：把多个小操作合并成一个

算子融合，说白了就是把几个连续的、可以合并的算子，打包成一个。为什么要这么做？因为每次调用一个算子，都有额外的调度开销。你想想看，如果能把三个算子合并成一个，那调度开销就直接省了三分之二。

3.1.1 Conv+BN+ReLU 融合

这是最经典的融合案例。在推理阶段，BatchNorm 其实就是一个线性变换：y = γ * (x - μ) / σ + β。而 Conv 也是一个线性变换。两个线性变换叠加，本质上还是一个线性变换。

我建议你在部署前，手动把 BN 的参数融合到 Conv 的权重里。具体做法是这样的：

// 伪代码：Conv+BN 融合
// 原始 Conv 权重: W_conv, b_conv
// BN 参数: gamma, beta, mean, var, epsilon

// 计算融合后的权重
W_fused = (gamma / sqrt(var + epsilon)) * W_conv

// 计算融合后的偏置
b_fused = gamma * (b_conv - mean) / sqrt(var + epsilon) + beta

// 然后直接使用 W_fused 和 b_fused 做 Conv，后面再接 ReLU

嗯，这里要注意：融合后的 Conv 层，后面直接跟 ReLU 就行。因为 BN 已经被吃掉了。

核心收益：减少一次 BN 的访存和计算，同时减少一次 kernel launch 开销。在端侧设备上，这个优化通常能带来 10%-20% 的加速。

3.1.2 其他常见的融合模式

除了 Conv+BN+ReLU，还有几种常见的融合模式：

  • Conv+ReLU 融合：直接把 ReLU 的激活函数嵌入到 Conv 的计算循环里，省掉一次额外的访存。
  • Add+ReLU 融合：在残差网络中很常见，把两个分支相加后直接激活。
  • Concat+Conv 融合：有些框架会把 Concat 和后面的 Conv 合并，减少中间结果的写回。

我在项目中遇到过一个问题：融合后的算子，数值精度可能会略有变化。因为 BN 的融合涉及到浮点运算的舍入误差。不过对于推理来说，这点误差完全可以接受。

3.2 常量折叠：把能算的先算好

常量折叠，说白了就是把那些输入全是常量的子图，提前计算好。比如一个 Add 操作，两个输入都是常量，那完全可以在编译阶段就把结果算出来，运行时直接拿结果就行。

你想想看，模型里其实有很多这样的机会：

  • Shape 相关的操作：比如 Reshape、Transpose 的输入 shape 如果是固定的，那输出 shape 也是固定的。
  • 初始化相关的操作：比如 Fill、Constant 这些，直接就是常量。
  • 数学运算：比如 Mul(2.0, 3.0)，直接算成 6.0。

我的习惯：在模型导出时，我会先用 ONNX 的常量折叠 pass 跑一遍。这样导出的模型就已经干净很多了。然后再做后续的量化或者剪枝。

我曾经踩过一个坑：常量折叠把一些动态 shape 的信息给丢了。比如某个 Reshape 的 shape 参数，在训练时是固定的，但推理时可能会变。如果强行折叠了，推理时就会报错。所以，常量折叠一定要确认 shape 是真正静态的。

3.3 死节点消除：把没用的东西扔掉

死节点消除，说白了就是删除那些对最终输出没有贡献的算子。这些算子可能是：

  • 训练时留下的：比如 dropout 层、frozen 的 BN 层。
  • 调试用的：比如打印中间结果的 Print 节点。
  • 冗余的：比如两个连续的 Identity 操作。

我建议你在部署前，先跑一遍死节点消除。具体做法是：

  1. 标记所有输出节点：模型的最终输出，以及被外部引用的中间节点。
  2. 反向遍历：从输出节点开始，反向遍历所有依赖的节点。
  3. 删除未标记的节点：那些没有被遍历到的节点，就是死节点，直接删掉。

注意：有些框架的图优化 pass 可能会误删一些看似无用、但实际上有 side effect 的节点。比如 Assert 节点，虽然对输出没贡献，但会检查输入合法性。删掉后可能导致运行时错误。

我记得有一次，我在优化一个 NLP 模型时，发现模型里有一堆 Shape 和 Gather 节点。这些节点在训练时用来动态计算序列长度，但推理时序列长度是固定的。跑完死节点消除后，模型大小直接减少了 30%。

3.4 图优化的整体流程

下面这张图展示了图优化的整体流程。你可以看到，算子融合、常量折叠、死节点消除这三个步骤是相互配合的。

图优化整体流程 原始计算图 常量折叠 死节点消除 算子融合 优化后计算图 优化顺序建议 1. 先做常量折叠 2. 再做死节点消除 3. 最后做算子融合 预期收益 模型大小减少 20-50% 推理速度提升 10-30% 内存占用降低 15-40%

从图中你可以看到，常量折叠和死节点消除是基础，它们把图清理干净后，算子融合才能发挥最大效果。我建议你按照这个顺序来优化。

3.5 实战中的注意事项

最后，分享几个我在实战中总结的经验：

  • 不要过度优化：有些融合可能会破坏后续的量化机会。比如 Conv+BN 融合后，BN 的参数就没了，后续的量化可能就做不了了。
  • 验证数值一致性：每次优化后，都要跑一遍测试集，确保输出和优化前一致。我一般会对比 top-1 和 top-5 准确率，误差在 0.1% 以内就算通过。
  • 考虑硬件特性：不同的端侧芯片，对算子融合的支持不一样。比如有些芯片的 NPU 原生支持 Conv+ReLU 融合，那你就没必要手动做了。

我的建议：如果你用的是 ONNX Runtime 或者 TensorRT，它们自带的图优化 pass 已经很强了。你只需要确保模型导出正确，然后交给它们就行。但如果你是自己写推理引擎，那这些优化就得自己实现了。

嗯，图优化这部分内容就到这里。记住一句话：优化不是目的，让模型在端侧跑得快、跑得稳才是目的。别为了优化而优化，有时候保持原样反而更好。

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