趋近智
深度学习 (deep learning)模型中的并非所有操作都依赖于输入数据。许多计算涉及固定参数 (parameter)、超参数 (hyperparameter)或配置值,这些值在每次推理 (inference)运行时保持不变。当编译器遇到这些静态子图时,它提供了一个机会,可以在编译期间执行一次计算,而不是为每个输入样本重复执行,从而降低运行时负载。这种优化方法称为常量折叠。
常量折叠分析计算图,以找出所有输入都是静态已知值的操作。编译器评估这些操作,并将原始子图替换为包含预计算结果的单个节点。此过程与常量传播配合使用,新计算出的常量被传递到后续操作,可能由此引发更多折叠机会的连锁反应。
常量折叠的主要目标是缩小计算图。在典型的神经网络 (neural network)中,常量以多种形式出现:
考虑一个场景,模型包含一个源自两个常量的缩放因子。在Python中,这可能看起来像 x * (a / b),其中 a 和 b 是在模型配置中定义的标量。如果没有优化,执行引擎在模型每次运行时都会执行 a / b 除法。常量折叠会检测到 a 和 b 是不可变的,在编译时计算商 c = a / b,然后重写图以执行 x * c。
以下图示说明了编译器如何识别并折叠静态子图。
一个图的转换,其中静态加法操作被预计算并替换为单个常量节点。
为实现常量折叠,编译器必须遍历中间表示(IR),并判断哪些节点是静态求值的候选者。这通常通过数据流分析方法实现。
编译器遍历图,通常按拓扑顺序。对于每个节点,它检查其传入边(操作数)的属性。如果一个节点的所有操作数都是常量节点,则该节点被认为是“可折叠”的。
如果一个节点是可折叠的,编译器会调用解释器或轻量级执行引擎来计算结果。这个解释器的复杂程度各不相同。对于简单算术,编译器可以使用主机CPU的本地指令。对于复杂张量操作,编译器可能需要链接到算子的参考实现,以确保编译时结果与目标硬件的运行时行为一致。
这种优化的优势源于其递归特性。当一个节点被折叠后,它实际上成为图中的一个新常量。这个新常量可能正是允许下游节点也进行折叠的缺失部分。
设想一个表示以下计算的图:
Sqrt 节点以 (常量)作为输入。它计算 。 Sqrt 节点被替换为包含 的常量节点。Add 节点最初接收 Sqrt 的输出和常量 。现在,它看到两个常量: 和 。Add 节点现在可折叠。编译器计算 。这种级联效应要求优化过程迭代运行或使用工作列表算法,直到图达到稳定状态,不再能进行任何简化。
常量折叠的一个实用且影响大的应用出现在推理 (inference)期间的批归一化(BatchNorm)层中。在训练期间,BatchNorm 会减去批量均值并除以批量标准差。然而,在推理期间,这些统计数据是固定的移动平均值。
批归一化的标准推理公式如下:
此处,(均值)、(方差)、(缩放)和 (偏移)都是学习参数 (parameter), 是一个配置常量。
朴素执行会在运行时进行减法、加法、平方根、除法、乘法和另一个加法。然而,由于 在推理时都是常量,编译器可以将它们折叠成两个不同的因子:一个乘法缩放因子 和一个偏置 (bias)项 。
运行时操作简化为线性变换 。如果批归一化层紧随一个卷积层之后,这些新常量通常可以直接折叠到卷积本身的权重 (weight)中,从而完全从运行时图中移除批归一化算子。
常量折叠不限于浮点张量数据。它大量用于与张量形状相关的整数算术。现代动态框架通常会生成复杂的子图,仅仅是为了计算 Reshape 或 Transpose 操作后张量的形状。
例如,将张量从 [Batch, Channels, Height, Width] 展平为 [Batch, Features] 涉及对维度大小的算术运算。如果输入图像大小是固定的(例如 ),编译器可以预计算展平向量 (vector)的精确大小。这消除了推理 (inference)内核中的整数算术指令,并使内存分配器能够在模型运行前准确知道需要多少缓冲区空间。
常量折叠通常有益,但也存在编译器工程师必须处理的极端情况。
代码大小膨胀: 如果常量计算生成的张量显著大于原始输入(例如,在小常量上使用 np.tile 或 torch.repeat),对其进行折叠可能会在编译后的二进制文件中存储一个巨大的数组。这会增加二进制文件大小和内存占用。如果结果常量超过某个大小阈值,编译器通常会使用启发式方法阻止折叠。
精度差异: 执行编译的机器可能与目标加速器有不同的浮点精度或舍入行为。在x86 CPU上预计算的值,与在嵌入 (embedding)式NPU或GPU上计算得到的结果,可能存在位级差异。对于数值敏感的模型,这种差异可能有效改变模型的准确性。
通过积极识别并处理静态图组件,常量折叠作为清理阶段,简化结构以暴露内存布局和循环调度方面的更多优化机会。
这部分内容有帮助吗?
© 2026 ApX Machine LearningAI伦理与透明度•