3、常见融合模式（一）：Conv+BatchNorm+ReLU 融合，原理与实现

说到算子融合，Conv+BatchNorm+ReLU 这个组合，绝对是实战中最常见、也最值得优先优化的模式。我刚开始做模型部署时，第一个优化的就是它。你想想看，一个典型的卷积神经网络里，这种「卷积+归一化+激活」的三件套，几乎占了计算量的 60% 以上。

为什么要融合？说白了，就是减少访存和 kernel launch 的开销。原本三个算子要读写三次显存，启动三次 GPU 内核。融合后，一次读写，一次启动。在训练吞吐上，这个收益非常可观。

3.1 融合的数学原理

我们先从公式入手。Conv 的数学表达很简单：

y = W * x + b

其中 W 是权重，x 是输入，b 是偏置。

BatchNorm 的推理阶段公式是：

z = γ * (y - μ) / √(σ² + ε) + β

这里 γ 是缩放因子，β 是偏移，μ 是均值，σ² 是方差，ε 是个小常数防止除零。

ReLU 就更直接了：

out = max(0, z)

好，现在我们把 Conv 的输出 y 代入 BatchNorm，再代入 ReLU。整个过程可以合并成一个新的卷积操作：

out = max(0,  (γ / √(σ² + ε)) * (W * x + b - μ) + β )

整理一下，令：

W' = (γ / √(σ² + ε)) * W
b' = (γ / √(σ² + ε)) * (b - μ) + β

那么融合后的算子就是：

out = max(0, W' * x + b')

看到了吗？一个带 ReLU 的 Conv 就搞定了。权重和偏置在融合阶段提前算好，训练时直接使用。

核心要点： 融合后的 Conv 层，权重做了缩放，偏置做了平移。这个变换在训练开始前做一次就行，不影响反向传播的梯度计算。

3.2 融合的收益分析

我在项目中遇到过这样一个场景：一个 ResNet-50 模型，单卡训练时，Conv+BN+ReLU 这三个算子占了 70% 的 GPU 时间。融合后，训练吞吐提升了 22%。

具体收益来自三个方面：

  • 减少 kernel 启动次数：从 3 次降到 1 次，省掉了 GPU 的调度开销
  • 减少显存读写：中间结果 y 和 z 不再需要写回显存，直接流水线处理
  • 减少显存占用：不需要为中间激活值分配额外空间，可以训练更大的 batch size

我习惯用一个简单的表格来对比：

指标 融合前 融合后 提升
Kernel 启动次数 3 1 66%
显存读写量 3N N 66%
中间激活显存 2N 0 100%

N 是特征图大小。你想想看，对于大模型来说，这个节省非常可观。

3.3 融合的实现细节

实现上，我建议分两步走：

  1. 权重和偏置的预计算：在训练开始前，把 BN 的参数融合进 Conv 的权重和偏置
  2. 替换网络中的算子：用融合后的 Conv+ReLU 替换原来的三个算子

下面是一个 PyTorch 风格的伪代码，展示了融合过程：

def fuse_conv_bn_relu(conv, bn):
    # 提取 BN 参数
    gamma = bn.weight.data
    beta = bn.bias.data
    mean = bn.running_mean
    var = bn.running_var
    eps = bn.eps

    # 计算缩放因子
    std = torch.sqrt(var + eps)
    scale = gamma / std

    # 融合权重
    fused_weight = conv.weight.data * scale.view(-1, 1, 1, 1)

    # 融合偏置
    if conv.bias is not None:
        fused_bias = (conv.bias.data - mean) * scale + beta
    else:
        fused_bias = -mean * scale + beta

    # 创建新的 Conv 层，带上 ReLU
    fused_conv = torch.nn.Conv2d(
        in_channels=conv.in_channels,
        out_channels=conv.out_channels,
        kernel_size=conv.kernel_size,
        stride=conv.stride,
        padding=conv.padding,
        bias=True
    )
    fused_conv.weight.data = fused_weight
    fused_conv.bias.data = fused_bias

    # 返回融合后的 Conv + ReLU
    return torch.nn.Sequential(fused_conv, torch.nn.ReLU(inplace=True))
小技巧： 我习惯在融合后做一次数值校验。拿一个随机输入，分别跑融合前和融合后的网络，对比输出。误差应该在 1e-5 以内才算合格。

3.4 避坑指南

我曾经踩过一个坑：BN 的 training 模式和 eval 模式混用。融合时，BN 必须处于 eval 模式，使用 running_mean 和 running_var。如果在 training 模式下融合，用的是 batch 的统计量，推理时就会出问题。

另一个常见问题是：融合后的 Conv 不支持 BN 的反向传播。因为 BN 的参数已经被吸收进 Conv 的权重里了。所以，融合通常只在推理阶段做，或者在训练的最后几个 epoch 做一次，然后固定住。

注意： 如果你在训练过程中做融合，一定要确保后续不再更新 BN 的参数。否则，融合后的权重和实际 BN 参数会不一致，导致精度下降。

3.5 融合后的网络结构

为了让你更直观地理解，我画了一张流程图：

Conv2D BatchNorm ReLU Fused Conv+ReLU (权重已融合 BN 参数) 融合 融合前：3 次 kernel 启动，3 次显存读写 融合后：1 次 kernel 启动，1 次显存读写 训练吞吐提升 15%~25%

嗯，这张图很直观。左边是融合前的三个独立算子，右边是融合后的一个算子。中间的虚线箭头就是我们的融合操作。

3.6 实战中的注意事项

最后，分享几个我在实战中总结的经验：

  • 不要盲目融合所有层：有些层的 BN 参数差异很大，融合后精度会掉。我一般只融合训练稳定的层。
  • 注意 inplace 操作：ReLU 的 inplace=True 可以节省显存，但融合后要确保不会影响梯度计算。
  • 多卡训练时小心：如果用了 SyncBN，融合前要确保所有卡的 BN 参数已经同步。

说白了，Conv+BN+ReLU 融合是算子优化的入门必修课。它原理简单，收益明显，实现也不复杂。我建议你拿到一个新模型，第一件事就是看看能不能做这个融合。往往能带来立竿见影的效果。

想要更多资料，联系可可

← 上一章 📖 返回目录 下一章 →