深度学习模型中的并非所有操作都依赖于输入数据。许多计算涉及固定参数、超参数或配置值,这些值在每次推理运行时保持不变。当编译器遇到这些静态子图时,它提供了一个机会,可以在编译期间执行一次计算,而不是为每个输入样本重复执行,从而降低运行时负载。这种优化方法称为常量折叠。常量折叠分析计算图,以找出所有输入都是静态已知值的操作。编译器评估这些操作,并将原始子图替换为包含预计算结果的单个节点。此过程与常量传播配合使用,新计算出的常量被传递到后续操作,可能由此引发更多折叠机会的连锁反应。静态求值的机制常量折叠的主要目标是缩小计算图。在典型的神经网络中,常量以多种形式出现:模型权重和偏置: 这些在训练后固定。超参数: 例如归一化epsilon、裁剪阈值或Leaky ReLU斜率等值。形状算术: 用于确定reshape或slice操作的张量维度的计算。考虑一个场景,模型包含一个源自两个常量的缩放因子。在Python中,这可能看起来像 x * (a / b),其中 a 和 b 是在模型配置中定义的标量。如果没有优化,执行引擎在模型每次运行时都会执行 a / b 除法。常量折叠会检测到 a 和 b 是不可变的,在编译时计算商 c = a / b,然后重写图以执行 x * c。以下图示说明了编译器如何识别并折叠静态子图。digraph G { rankdir=LR; node [shape=box, style=filled, fontname="Helvetica", fontsize=12]; splines=ortho; subgraph cluster_0 { label="优化前"; style=dashed; color="#adb5bd"; fontcolor="#868e96"; node [color="#ced4da", fillcolor="#f8f9fa"]; Input [label="输入张量", fillcolor="#e7f5ff", color="#74c0fc"]; W1 [label="常量: 5.0"]; W2 [label="常量: 2.0"]; node [color="#bac8ff", fillcolor="#dbe4ff"]; Add [label="加法"]; Mul [label="乘法"]; Input -> Mul; W1 -> Add; W2 -> Add; Add -> Mul; } subgraph cluster_1 { label="优化后"; style=dashed; color="#adb5bd"; fontcolor="#868e96"; node [color="#ced4da", fillcolor="#f8f9fa"]; Input2 [label="输入张量", fillcolor="#e7f5ff", color="#74c0fc"]; Folded [label="折叠常量: 7.0", fillcolor="#d8f5a2", color="#94d82d"]; node [color="#bac8ff", fillcolor="#dbe4ff"]; Mul2 [label="乘法"]; Input2 -> Mul2; Folded -> Mul2; } }一个图的转换,其中静态加法操作被预计算并替换为单个常量节点。识别可折叠节点为实现常量折叠,编译器必须遍历中间表示(IR),并判断哪些节点是静态求值的候选者。这通常通过数据流分析方法实现。编译器遍历图,通常按拓扑顺序。对于每个节点,它检查其传入边(操作数)的属性。如果一个节点的所有操作数都是常量节点,则该节点被认为是“可折叠”的。如果一个节点是可折叠的,编译器会调用解释器或轻量级执行引擎来计算结果。这个解释器的复杂程度各不相同。对于简单算术,编译器可以使用主机CPU的本地指令。对于复杂张量操作,编译器可能需要链接到算子的参考实现,以确保编译时结果与目标硬件的运行时行为一致。传播与级联这种优化的优势源于其递归特性。当一个节点被折叠后,它实际上成为图中的一个新常量。这个新常量可能正是允许下游节点也进行折叠的缺失部分。设想一个表示以下计算的图: $$y = x + \sqrt{100} + 5$$初始状态: 图中包含 $x$ 的输入、常量 $100$ 和常量 $5$。第一次处理: 编译器看到 Sqrt 节点以 $100$(常量)作为输入。它计算 $\sqrt{100} = 10$。 Sqrt 节点被替换为包含 $10$ 的常量节点。传播: 下游的 Add 节点最初接收 Sqrt 的输出和常量 $5$。现在,它看到两个常量:$10$ 和 $5$。第二次处理: Add 节点现在可折叠。编译器计算 $10 + 5 = 15$。最终状态: 图被简化为 $y = x + 15$。这种级联效应要求优化过程迭代运行或使用工作列表算法,直到图达到稳定状态,不再能进行任何简化。批归一化中的优化常量折叠的一个实用且影响大的应用出现在推理期间的批归一化(BatchNorm)层中。在训练期间,BatchNorm 会减去批量均值并除以批量标准差。然而,在推理期间,这些统计数据是固定的移动平均值。批归一化的标准推理公式如下:$$y = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \gamma + \beta$$此处,$\mu$(均值)、$\sigma^2$(方差)、$\gamma$(缩放)和 $\beta$(偏移)都是学习参数,$\epsilon$ 是一个配置常量。朴素执行会在运行时进行减法、加法、平方根、除法、乘法和另一个加法。然而,由于 $\mu, \sigma, \epsilon, \gamma, \beta$ 在推理时都是常量,编译器可以将它们折叠成两个不同的因子:一个乘法缩放因子 $w_{new}$ 和一个偏置项 $b_{new}$。$$w_{new} = \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}}$$ $$b_{new} = \beta - \frac{\mu \gamma}{\sqrt{\sigma^2 + \epsilon}}$$运行时操作简化为线性变换 $y = w_{new}x + b_{new}$。如果批归一化层紧随一个卷积层之后,这些新常量通常可以直接折叠到卷积本身的权重中,从而完全从运行时图中移除批归一化算子。形状推断与静态切片常量折叠不限于浮点张量数据。它大量用于与张量形状相关的整数算术。现代动态框架通常会生成复杂的子图,仅仅是为了计算 Reshape 或 Transpose 操作后张量的形状。例如,将张量从 [Batch, Channels, Height, Width] 展平为 [Batch, Features] 涉及对维度大小的算术运算。如果输入图像大小是固定的(例如 $224 \times 224$),编译器可以预计算展平向量的精确大小。这消除了推理内核中的整数算术指令,并使内存分配器能够在模型运行前准确知道需要多少缓冲区空间。权衡与局限常量折叠通常有益,但也存在编译器工程师必须处理的极端情况。代码大小膨胀: 如果常量计算生成的张量显著大于原始输入(例如,在小常量上使用 np.tile 或 torch.repeat),对其进行折叠可能会在编译后的二进制文件中存储一个巨大的数组。这会增加二进制文件大小和内存占用。如果结果常量超过某个大小阈值,编译器通常会使用启发式方法阻止折叠。精度差异: 执行编译的机器可能与目标加速器有不同的浮点精度或舍入行为。在x86 CPU上预计算的值,与在嵌入式NPU或GPU上计算得到的结果,可能存在位级差异。对于数值敏感的模型,这种差异可能有效改变模型的准确性。通过积极识别并处理静态图组件,常量折叠作为清理阶段,简化结构以暴露内存布局和循环调度方面的更多优化机会。