2、计算图优化：图融合、常量折叠、死代码消除
好，咱们接着聊计算图优化。说实话，这是嵌入式AI部署里最「润物细无声」的一环。你模型结构写得再花哨，推理框架在底层给你做一波优化，性能可能直接翻倍。我个人习惯把计算图优化比作「给代码做减法」——去掉没用的，合并能合并的，提前算好能算的。
今天咱们重点聊三个经典优化：图融合、常量折叠、死代码消除。这三个东西，说白了就是让计算图更「瘦」、更「快」、更「聪明」。
2.1 图融合：把多个算子捏成一个
图融合，英文叫 Operator Fusion。什么意思呢？就是把计算图中连续的几个小算子，合并成一个大算子。为什么要这么做？因为每个算子在执行时，都有启动开销（kernel launch overhead）。你想想看，如果每个小算子都单独启动一次GPU或NPU，那光调度时间就够喝一壶的。
我在项目中遇到过最典型的例子，就是 Conv + BatchNorm + ReLU 的融合。这三个操作在推理时完全可以合并成一个算子。BN 的公式可以吸收进 Conv 的权重里，ReLU 直接跟在后面。这样一来，原本三次内存读写变成了一次，速度提升非常明显。
核心思路：减少 kernel launch 次数，减少中间张量的内存读写。
举个例子，假设你有这么一段计算图：
// 原始计算图：Conv -> BN -> ReLU
y = conv(x, w, b)
z = bn(y, gamma, beta, mean, var)
out = relu(z)融合之后，变成：
// 融合后：ConvBNReLU
w_fused = gamma * w / sqrt(var + eps)
b_fused = (b - mean) * gamma / sqrt(var + eps) + beta
out = relu(conv(x, w_fused, b_fused))嗯，这里要注意：融合后的权重是离线算好的，推理时直接加载。所以 BN 的参数在推理阶段其实已经「消失」了。这也是为什么很多部署框架要求你先把 BN 层 freeze 掉。
我的经验：在 TensorRT 里，Conv+BN+ReLU 的融合是自动做的。但在一些轻量级框架里，你得手动写融合 pass。我建议你在模型导出前，先用 torch 的 torch.nn.utils.fuse_conv_bn_eval 把 BN 融合进 Conv，这样导出后的图更干净。
2.2 常量折叠：能提前算的，绝不拖到运行时
常量折叠，Constant Folding。这个名字很形象——把计算图中那些输入全是常量的子图，提前算出来，替换成常量节点。
为什么会存在这种机会？因为很多时候，模型里会有一些固定的数学运算，比如 shape 计算、padding 计算、或者某些固定的 mask 生成。这些操作在推理时输入是不变的，那何必每次推理都算一遍？
我曾经调试过一个模型，发现每次推理都要花 2ms 去生成一个 attention mask。那个 mask 其实只跟序列长度有关，而序列长度在部署时是固定的。我直接把 mask 生成逻辑用常量折叠优化掉，推理时间直接降了 1.8ms。你看，这种优化有时候比调算子本身还管用。
举个例子：
// 原始计算图
shape = shape_of(input) // 假设 input shape 固定为 [1, 3, 224, 224]
size = prod(shape) // 1*3*224*224 = 150528
scale = mul(size, 0.5) // 75264
// 常量折叠后
scale = 75264 // 直接替换成常量节点注意：常量折叠不是万能的。如果输入 shape 是动态的，那这些常量折叠就做不了。所以部署时尽量固定输入尺寸，这样优化空间更大。
2.3 死代码消除：删掉那些「没用」的节点
死代码消除，Dead Code Elimination。这个其实很好理解——计算图里有些节点，它的输出没有被任何后续节点使用，那它就是「死」的，可以直接删掉。
你可能会问：模型训练的时候怎么会产生死代码？其实很常见。比如你训练时用了多个 loss 分支，但部署时只保留主分支；或者你做了某些实验性的中间计算，但最终没用到。这些节点在训练图里是活的，但在推理图里就是死的。
我记得有一次，一个同事导出的 ONNX 模型有 200 多个节点，但实际推理只需要 80 个。剩下的全是训练时留下的辅助节点，比如梯度计算节点、中间统计节点。这些在推理时完全没用，删掉之后模型体积小了 60%，推理速度也快了 30%。
死代码消除的规则很简单：
- 从输出节点开始反向遍历
- 标记所有可达的节点
- 删除所有不可达的节点
说白了，就是「从终点倒着走，走不到的统统扔掉」。
避坑指南：我曾经遇到过一个问题——删掉一个看似「死」的节点后，模型推理结果变了。后来发现那个节点虽然输出没被用，但它有 side effect（比如更新了某个全局状态）。所以在做死代码消除时，一定要确认节点没有副作用。大部分框架的优化 pass 会考虑这一点，但自己手写时千万小心。
2.4 三种优化的优先级与配合
这三种优化不是孤立的，它们之间有依赖关系。我个人习惯的顺序是：
- 先做常量折叠——把能算的算出来，这样后面的图更干净
- 再做死代码消除——把常量折叠后可能产生的无用节点删掉
- 最后做图融合——在精简后的图上做算子合并
为什么是这个顺序？你想想看，常量折叠可能会产生新的常量节点，这些节点如果没被用到，就是死代码。先做死代码消除，可以避免在死节点上做无意义的融合。嗯，这个顺序是我踩过坑之后总结出来的。
下面是一个简单的对比表格：
| 优化类型 | 目标 | 典型收益 | 注意事项 |
|---|---|---|---|
| 图融合 | 减少 kernel 启动次数 | 推理速度提升 20%-50% | 注意算子语义是否等价 |
| 常量折叠 | 减少运行时计算量 | 减少 5%-15% 的计算量 | 输入 shape 必须固定 |
| 死代码消除 | 减少模型体积和计算量 | 模型体积减少 10%-60% | 注意 side effect |
2.5 实战建议
最后给几个实战建议：
- 用现成的工具：TensorRT、ONNX Runtime、OpenVINO 这些框架都内置了这些优化。你不需要自己实现，但需要理解它们做了什么，以便排查问题。
- 可视化检查：我习惯用 Netron 查看优化前后的计算图。一眼就能看出哪些节点被融合了，哪些被删掉了。
- 不要过度优化：有时候融合太激进，反而会导致数值精度问题。比如 Conv+BN 融合后，如果 BN 的 eps 很小，可能会引入数值误差。建议融合后做一次精度对比。
好了，这一章就聊到这儿。下一章咱们会深入聊聊内存优化和算子调度，那又是另一番天地了。