计算图优化：让模型跑得更快的四把手术刀

各位同学，今天我们来聊聊计算图优化。说实话，这玩意儿刚接触时觉得挺玄乎的，但干久了你会发现——它其实就是把模型的计算过程「捋顺」了。

我经常跟团队里的小朋友说：你写出来的模型，跟框架实际执行的，往往是两码事。框架会在背后偷偷做很多手脚，目的就一个——让你算得更快。

计算图优化，说白了就是四招：图融合、常量折叠、死代码消除、算子替换。咱们一个一个来拆。

1. 图融合：把零碎操作打包

先问个问题：为什么要把多个算子合并成一个？

你想想看，每个算子在GPU上执行，都有启动开销。哪怕是个极简单的加法，也得走一遍「分配显存→加载数据→执行→写回」的流程。算子越多，启动开销越大。

我遇到过最夸张的一个案例：某同学的模型里，一个卷积后面跟了四个独立的激活操作。每个激活都单独写成了一个算子。结果呢？GPU利用率不到30%。

图融合的核心思路：把连续的、可以合并的算子打包成一个。

举个例子，你写的是这样的计算图：

# 原始计算图
x → Conv2D → BatchNorm → ReLU → y

# 融合后的计算图
x → FusedConvBNReLU → y

嗯，这里要注意：不是所有算子都能随便融合。比如有些框架要求融合后的算子必须保持数值精度一致。我曾经踩过坑，把FP16的卷积和FP32的BN强行融合，结果精度掉了0.3个点。

2. 常量折叠：提前算好，别等运行时

这个其实很好理解。如果某个节点的输入全是常量，那为什么不在编译时就把它算出来？

我见过不少模型里，有人写了这样的代码：

# 原始计算图
const_a = [1, 2, 3, 4, 5]
const_b = [10, 20, 30, 40, 50]
c = const_a + const_b  # 这个加法每次推理都要算一遍

# 常量折叠后
c = [11, 22, 33, 44, 55]  # 直接替换成常量

你可能会说：「这谁看不出来啊？」但实际模型里，这种常量计算往往藏在很深的子图里。比如某个Shape操作后面跟了一堆Reshape、Slice，输入全是固定的——那这些操作完全可以在编译期完成。

3. 死代码消除：砍掉没用的分支

这个听起来有点吓人，但说白了就是：删掉那些对最终输出没有贡献的节点。

为什么会存在死代码？

• 调试时加的打印节点，忘了删
• 从别的模型里复制过来的子图，有些分支没用上
• 条件控制流里，某些分支永远不会被执行

我曾经接手过一个老项目，模型文件有200MB。一分析，好家伙，里面挂了十几个没用的Identity节点和Print节点。砍掉之后，模型直接瘦身到180MB，推理速度还快了5%。

4. 算子替换：用更快的实现替代

这个就更有意思了。同一个数学运算，可以用不同的算子来实现，但速度天差地别。

举个例子：

原始算子	替换算子	加速比（实测）
Conv2D (3x3, stride=1)	Winograd Conv2D	1.5x ~ 2x
BatchNorm (推理时)	Scale + Bias	3x ~ 5x
Softmax (大维度)	Online Softmax	2x ~ 3x
LayerNorm	Fused LayerNorm	1.8x ~ 2.5x

我印象最深的是BatchNorm的替换。很多同学不知道，推理时的BatchNorm其实就是一个线性变换——完全可以用一个Scale加一个Bias来代替。去掉那些除方差、加epsilon的操作，速度能快好几倍。

算子替换的另一个经典场景是大算子拆小。比如一个大矩阵乘法，显存放不下怎么办？拆成多个小矩阵乘法，分块计算。虽然总计算量没变，但显存压力小了，反而能跑更大的batch size。

一张图看懂计算图优化

下面这张图，是我自己总结的优化流程。你照着这个思路走，基本不会漏掉什么：

避坑指南

最后，分享几个我踩过的坑：

• 别盲目追求融合。有些框架的融合算子实现得并不好，反而比分开执行还慢。我建议你每次融合后都做一次benchmark。
• 常量折叠要注意精度。浮点数常量折叠时，不同框架的舍入策略可能不一样。跨平台部署时要格外小心。
• 死代码消除要留个心眼。有些节点虽然输出没被用到，但可能有副作用（比如更新全局计数器）。删之前先确认。
• 算子替换要考虑硬件。同一个算子，在NVIDIA GPU上快，在AMD GPU上可能慢。别拿一套优化方案到处套。

好了，计算图优化这块就聊到这儿。这四招你掌握了，大部分模型的推理速度都能提一截。下次遇到模型跑得慢，别急着换硬件——先看看计算图，说不定动动刀子就解决了。

想要更多资料，联系可可