图优化入门：为什么需要图优化

各位同学，今天咱们聊聊图优化。说实话，我刚入行那会儿，觉得模型写好了能跑就行，图优化？那是锦上添花的事。直到有一次，我在一个边缘设备上部署模型，推理速度慢得让人抓狂，显存直接爆掉。嗯，从那以后，我再也不敢小看图优化了。

为什么需要图优化？

你想想看，我们写模型的时候，用的都是PyTorch、TensorFlow这些框架。框架为了方便我们开发，会把一个计算拆成很多细小的算子。比如一个简单的卷积+激活+池化，框架可能会生成十几个小算子。每个算子都要单独调度、分配显存、启动内核。这就像你搬家，明明可以一卡车拉走，非要分十趟跑——效率能高吗？

图优化要做的，就是把这些零碎的小算子合并成大算子，把没用的计算剪掉，把能提前算的提前算好。说白了，就是给计算图做一次「瘦身」和「重组」。

图优化的目标

图优化的目标其实就两个：加速和省显存。我分别说说。

加速

加速怎么来的？主要有三个途径：

• 减少内核启动次数：每个算子启动一次GPU内核都有开销。合并算子后，启动次数大幅减少。我在项目中遇到过，一个MobileNetV3经过算子融合后，内核启动次数从300多次降到了80多次，推理时间直接砍半。
• 减少内存读写：算子之间如果有中间结果，原来要写到显存再读出来。融合后，中间结果直接留在寄存器或共享内存里，省掉了大量的显存带宽消耗。
• 消除冗余计算：有些计算是重复的，比如两个分支都做了同样的reshape。图优化可以识别并消除这些冗余。

省显存

省显存这块，我印象特别深。有一次部署一个BERT模型，显存死活不够用。后来做了图优化，把一些不需要保留的中间张量及时释放掉，显存占用直接降了40%。

省显存的主要手段：

• 内存复用：两个算子的输入输出形状相同，可以共用同一块显存。
• 原地操作：有些算子可以原地修改输入，不需要额外分配输出。
• 消除中间张量：融合后的算子不需要保存中间结果，显存自然就省下来了。

图优化的基本流程

图优化不是一步到位的，它有一套标准流程。我习惯把它分成四个阶段：

1. 图解析：把模型的计算图读进来，构建成内部表示。这一步要搞清楚每个算子的输入输出、属性、依赖关系。
2. 图分析：遍历整个图，收集信息。比如哪些算子可以融合、哪些张量生命周期短、哪些分支是死代码。这一步是优化的基础，信息越全，优化效果越好。
3. 图变换：根据分析结果，对图进行修改。包括算子融合、常量折叠、死代码消除、内存优化等。这一步要小心，改错了模型就废了。
4. 图验证：变换后的图要和原图在数学上等价。我一般会跑几个测试用例，对比输出结果。误差在浮点精度范围内就算通过。

这里我画个简单的流程图，方便你理解：

原始计算图
    ↓
图解析 → 构建内部表示
    ↓
图分析 → 收集依赖、形状、生命周期等信息
    ↓
图变换 → 算子融合、常量折叠、死代码消除等
    ↓
图验证 → 等价性检查
    ↓
优化后的计算图

一个简单的例子

咱们看个具体的例子。假设你有这么一段计算：

# 原始计算
def original(x):
    a = conv2d(x, w1)
    b = relu(a)
    c = batch_norm(b)
    d = relu(c)
    return d

这个计算图里有四个算子：Conv2D、ReLU、BatchNorm、ReLU。每个算子都要单独调度，中间结果a、b、c都要写到显存里。

经过图优化后，可以变成这样：

# 优化后的计算
def fused(x):
    return fused_conv_bn_relu(x, w1)

一个融合算子搞定。中间结果a、b、c都不需要显式保存，直接在寄存器里流转。内核启动次数从4次降到1次，显存读写也大幅减少。

你想想看，一个简单的四算子链就能省这么多，一个几百个算子的模型优化空间有多大？

常见的图优化技术

我整理了几种最常用的图优化技术，你在实际项目中大概率都会用到：

优化技术	说明	效果
算子融合	将多个连续算子合并为一个	减少内核启动和内存读写
常量折叠	提前计算常量表达式	减少运行时计算量
死代码消除	删除不会被用到的计算节点	减少无用计算
内存复用	让不同张量共享同一块内存	降低峰值显存
算子替换	用更高效的实现替换原算子	提升计算效率

好了，这一章就到这里。下一章咱们会深入讲算子融合的具体实现，包括Conv+BN+ReLU的融合细节。到时候我会拿实际代码一步步演示，保证你学完就能用。