趋近智
机器学习模型通常包含在最终执行中不起作用的操作。这些多余的节点可能源于代码中遗留的调试语句、从不执行的条件分支,或训练逻辑中在推理时无关紧要的残留。在编译器设计中,移除这些不必要操作被称为死代码消除(DCE)。
虽然标准软件编译器移除无法访问的代码以减小二进制文件大小,但机器学习编译器执行 DCE 主要为了节省计算资源和内存带宽。计算图中的每个操作都会耗费时间和能量。如果张量操作的结果从未被模型的输出或任何后续节点读取,那么计算它就是资源浪费。
神经网络中的死代码通常源于三个方面。首先是用户定义。在试验模型架构时,您可能会定义用于可视化或特征提取的辅助分支,它们并非主要预测路径的一部分。如果这些分支被捕获到中间表示(IR)中,但其输出未被请求,它们就会成为冗余部分。
第二个来源是不同优化阶段的关联。像常量折叠这样的技术常常会作为副产品产生死代码。如果一个条件语句依赖于编译器确定实际为常量的值(例如 if use_dropout:,其中 use_dropout 为 False),那么与真条件关联的整个分支将变得不可达。编译器必须运行 DCE 来清除折叠操作遗留的无用部分。
第三个来源是从训练到推理的转变。训练图包含反向传播操作、梯度计算和损失函数。当编译模型以进行部署时,反向图是不必要的。虽然大多数框架会明确分离这些,但复杂的自定义层可能会意外保留对训练专用张量的引用。
为了识别哪些节点可以安全移除,编译器会进行活跃性或可达性分析。此过程将计算图视为一个有向无环图(DAG),其中边表示数据依赖关系。
该算法通常按以下步骤运作:
考虑一个简化的计算图,其中一个旁支计算用于记录的标准差,但该值从未被函数返回。
一个数据流图,显示了活跃节点(绿色)与死节点(灰色)。计算均值的分支源自 ReLU 输出,但未连接到最终图输出,使其成为可消除的对象。
机器学习编译器中的一个情况是对具有副作用的操作的处理。在纯数据流图中,一个操作只有在其数据被使用时才有意义。然而,某些操作,例如 print() 语句或自定义日志操作符,不会产生被数学管道使用的张量输出,但会执行对用户或系统可见的操作。
如果编译器激进地移除所有没有数据后继的节点,它可能会移除这些日志操作。为了避免这种情况,IR 通常支持将特定节点标记为“非纯”或具有“副作用”。活跃性分析算法将这些节点视为辅助根。如果一个节点有副作用,它会被标记为活跃,而且重要的是,它的所有输入依赖也都会被标记为活跃。
死代码消除很少作为单次遍历运行。它通常是迭代优化循环的一部分。这是因为移除一个死节点可能会使其输入没有剩余消费者,从而导致它们也成为死节点。
考虑以下操作序列:
t1t2t3返回 =x+y=t1×2=ReLU(t2)t1在初步分析中,根是返回值 t1。
根据实现情况,编译器可能通过引用计数在一次反向遍历中识别 t3 和 t2,或者可能需要多次遍历,其中消除 t3 会使 t2 成为新的消除对象。
这种循环关系在与其他优化结合使用时特别有意义。例如:
DCE 的好处不仅限于跳过算术指令。最显著的收益常常来自内存管理。在上面的例子中,如果 t3 和 t2 被消除,编译器就不需要为这些张量分配中间缓冲区。
对于大型神经网络,中间激活图可以消耗数千兆字节的显存。通过仔细裁剪图,编译器减少了峰值内存占用,可能允许更大的批次大小或更大的模型适应设备。当为内存限制严格的边缘设备优化时,这尤其有意义。
在诸如 TVM 或 XLA(加速线性代数)等框架中,DCE 作用于该栈的特定中间表示(例如,Relay IR 或 HLO)。
检查此阶段之前和之后的 IR 时,您会看到更清晰的结构。
DCE 之前:
fn (%x: Tensor[(10, 10), float32]) {
%0 = add(%x, 1.0f);
%1 = multiply(%0, 2.0f); /* 死代码: %1 从未被使用 */
%2 = relu(%0);
%2
}
DCE 之后:
fn (%x: Tensor[(10, 10), float32]) {
%0 = add(%x, 1.0f);
%2 = relu(%0);
%2
}
请注意,%1 及其相关计算被完全移除。依赖 %0 保留,因为它为 %2 所需,而 %2 是函数返回值。
尽管理论上很简单,死代码消除是一个基础性阶段,它确保后续复杂转换的效率。它充当图优化阶段的垃圾回收器,确保硬件只执行为产生预期预测而绝对必需的计算。
这部分内容有帮助吗?
© 2026 ApX Machine Learning用心打造