2、计算图优化：图融合、常量折叠、死代码消除

好，咱们今天聊聊计算图优化。这玩意儿说白了，就是让模型跑得更快、更省内存的「魔法」。你想想看，一个模型从训练完到部署，中间其实有很多「赘肉」——多余的节点、重复的计算、可以提前算好的常量。把这些清理掉，推理速度能快不少。

我个人习惯把计算图优化分成三大类：图融合、常量折叠、死代码消除。这三板斧用好了，模型部署的底子就稳了。

2.1 图融合：把多个算子「捏」成一个

图融合，顾名思义，就是把多个连续的算子合并成一个。为什么要这么做？因为每个算子执行时，都有启动开销——GPU 要加载 kernel、读写中间结果。算子越多，这些开销越大。

我在项目中遇到过最典型的例子，就是 Conv + BN + ReLU 的融合。训练时 BN 层是必要的，但推理时 BN 的参数可以吸收到 Conv 的权重里。你看这个流程：

# 原始计算图：Conv -> BN -> ReLU
# 每个节点都要读写一次显存

# 融合后：ConvReLU（内部包含BN参数）
# 一次读写，一次kernel启动

嗯，这里要注意：不是所有算子都能随便融合。融合的前提是：算子之间没有分支，且数据类型、形状完全匹配。我曾经踩过一个坑，把 int8 的 Conv 和 float32 的 ReLU 强行融合，结果精度直接崩了。

常见的融合模式：

  • Conv + BN + ReLU → ConvReLU
  • Conv + Add + ReLU → ConvAddReLU
  • MatMul + Add → MatMulAdd
  • LayerNorm + ReLU → LayerNormReLU

你想想看，一个 50 层的 ResNet，每层都做一次融合，能省掉 50 次 kernel 启动和中间显存读写。这优化幅度，可不是小打小闹。

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

常量折叠，说白了就是「提前算好」。模型里有些节点的输入全是常量，比如 shape 计算、padding 参数、某些固定的数学运算。这些节点在推理时每次结果都一样，何必每次都算一遍？

举个例子：

# 原始计算图
a = Constant(3)
b = Constant(4)
c = Add(a, b)  # 每次推理都算 3+4

# 常量折叠后
c = Constant(7)  # 直接替换成常量

我建议你在做常量折叠时，重点关注这几类节点：

  • Shape 相关：比如获取 tensor 的维度、大小
  • 数学运算：两个常量相加、相乘、取整
  • 索引操作：从常量 tensor 中取固定位置的元素
  • 类型转换：int 转 float、float 转 int 等

小技巧：常量折叠最好在模型导出时就做一次。很多框架（比如 ONNX、TensorRT）在优化 pass 里已经内置了常量折叠，但有时会漏掉一些自定义算子。我习惯在导出后手动跑一遍常量折叠脚本，确保万无一失。

为什么会漏掉？因为有些框架的常量折叠只处理标准算子，遇到自定义算子就「罢工」了。这时候你得自己写一个 pass，把自定义算子的常量折叠逻辑加进去。

2.3 死代码消除：砍掉没用的节点

死代码消除，听着有点吓人，其实就是删掉那些「输出不被任何节点使用」的节点。模型在开发过程中，难免会留下一些「垃圾」——比如调试用的打印节点、废弃的分支、多余的 identity 操作。

我记得有一次，一个同事导出的模型里居然有 200 多个 identity 节点。这些节点什么都不做，就是把输入原样输出。但每个 identity 都会占用一次 kernel 启动和显存读写。砍掉之后，推理速度提升了 15%。

死代码消除的流程其实很简单：

  1. 从输出节点开始，反向遍历计算图
  2. 标记所有「可达」的节点
  3. 删除所有「不可达」的节点

注意：死代码消除不能乱删。有些节点虽然看起来「没用」，但可能是 side effect 节点（比如打印日志、更新计数器）。我建议只删除那些确定没有副作用的节点，比如 identity、cast、shape 等。

嗯，这里还要提一个常见的坑：循环中的死代码。有些节点在循环外部是死代码，但在循环内部是活的。比如：

# 错误示例：把循环外的常量节点删了
for i in range(10):
    x = x + Constant(1)  # Constant(1) 在循环外定义，但循环内使用

# 正确做法：保留 Constant(1)，因为它被循环内的节点引用

我曾经因为没处理好循环中的死代码，导致模型推理结果全错。从那以后，我每次做死代码消除都会先检查循环结构。

2.4 优化顺序与实战建议

这三类优化不是独立的，它们之间有依赖关系。我个人习惯的顺序是：

步骤 优化类型 说明
1 常量折叠 先算好所有常量，减少后续优化的节点数
2 死代码消除 砍掉常量折叠后产生的无用节点
3 图融合 在精简后的图上做融合，效果更好

你想想看，如果先做图融合，再常量折叠，那融合后的算子可能包含常量节点，折叠起来更麻烦。所以顺序很重要。

实战建议：

  • 用 ONNX 导出模型时，开启 optimize=True，它会自动做一部分优化
  • 用 TensorRT 时，它的优化器会自动做图融合和常量折叠
  • 如果自己写推理引擎，建议用 torch.fx 或 onnxruntime 的优化 pass
  • 每次优化后，用 onnx.checker 或 torch.onnx.export 验证模型正确性

好了，计算图优化这块就聊到这儿。说白了，就是帮模型「减减肥」——去掉赘肉、合并重复动作、提前算好该算的。下一章咱们聊聊更细粒度的算子优化，比如内存布局和指令选择。到时候再细说。

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