3、算子融合技术（上）：Conv+BatchNorm融合原理、Conv+ReLU融合、实战代码演示

各位同学，今天我们来聊聊算子融合。这玩意儿，说白了就是给模型做「瘦身手术」。我刚开始做部署的时候，总觉得模型训练完就万事大吉了。结果一上端侧设备，延迟直接爆炸。后来才发现，算子融合是优化里最立竿见影的手段之一。

为什么要做融合？你想想看，每个算子在执行时，都有启动开销、内存读写开销。把多个算子合并成一个，就能省掉中间那些不必要的搬运和调度。今天我们先讲两个最经典的组合：Conv+BatchNorm 和 Conv+ReLU。

Conv+BatchNorm 融合原理

这个组合在推理时几乎必做。为什么？因为 BatchNorm 在训练时有用，但推理时它的行为是确定的。说白了，推理时的 BN 就是一个线性变换：y = γ * (x - μ) / √(σ² + ε) + β。

而卷积本身也是线性变换：y = W * x + b。

两个线性变换叠在一起，完全可以合并成一个。我在项目中遇到过不少同学，训练完模型直接导出 ONNX，结果发现 BN 节点还在。嗯，这其实浪费了不少性能。

融合的数学推导很简单。假设卷积的输出是 z = W * x + b_conv，BN 的输入就是 z。那么 BN 的输出为：

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

把 z 代入：

y = (γ / √(σ² + ε)) * (W * x + b_conv) + (β - γ * μ / √(σ² + ε))
  = [ (γ / √(σ² + ε)) * W ] * x + [ (γ / √(σ² + ε)) * b_conv + (β - γ * μ / √(σ² + ε)) ]

看到了吗？新的卷积权重就是 W' = (γ / √(σ² + ε)) * W，新的偏置就是 b' = (γ / √(σ² + ε)) * b_conv + (β - γ * μ / √(σ² + ε))。

核心要点：融合后的卷积，权重和偏置都变了。但计算量没变，只是省掉了 BN 那一步的内存读写和 kernel 启动开销。

Conv+ReLU 融合

这个更简单。ReLU 就是 max(0, x)。很多硬件或推理框架（比如 TensorRT、ncnn）都支持在卷积计算完成后，直接对输出做 in-place 的 ReLU 操作。

为什么能融合？因为 ReLU 没有可学习参数，也不依赖 batch 统计量。它就是一个逐元素的非线性激活。你可以在卷积的最后一个循环里，直接把结果 clamp 到非负。

我曾经在优化一个 MobileNet 模型时，发现 Conv+BN+ReLU 三个算子串在一起。融合掉 BN 和 ReLU 后，推理速度提升了将近 30%。你想想看，这还只是三个算子的组合。

个人经验：有些框架（比如 ONNX Runtime）会自动做这些融合。但如果你用自定义算子或者手写推理代码，一定要手动实现。我建议你在导出模型时，就检查一下计算图里是否还有独立的 BN 节点。

实战代码演示

下面我用 PyTorch 演示一下 Conv+BatchNorm 的融合过程。注意，这里只做推理时的融合，训练时不要动。

import torch
import torch.nn as nn

def fuse_conv_bn(conv, bn):
    """
    融合 Conv2d + BatchNorm2d
    返回一个新的 Conv2d，权重和偏置已更新
    """
    # 确保是 eval 模式
    conv.eval()
    bn.eval()

    # 获取 BN 参数
    gamma = bn.weight.data
    beta = bn.bias.data
    mean = bn.running_mean
    var = bn.running_var
    eps = bn.eps

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

    # 更新卷积权重
    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

    # 创建新的卷积层
    fused_conv = nn.Conv2d(
        in_channels=conv.in_channels,
        out_channels=conv.out_channels,
        kernel_size=conv.kernel_size,
        stride=conv.stride,
        padding=conv.padding,
        dilation=conv.dilation,
        groups=conv.groups,
        bias=True
    )

    fused_conv.weight.data = fused_weight
    fused_conv.bias.data = fused_bias

    return fused_conv

# 使用示例
conv = nn.Conv2d(3, 16, 3, padding=1)
bn = nn.BatchNorm2d(16)
# 模拟训练后的状态
bn.running_mean = torch.randn(16)
bn.running_var = torch.rand(16) + 1.0

fused_conv = fuse_conv_bn(conv, bn)

# 验证结果是否一致
x = torch.randn(1, 3, 32, 32)
with torch.no_grad():
    out_original = bn(conv(x))
    out_fused = fused_conv(x)

print("最大误差:", (out_original - out_fused).abs().max().item())
# 应该接近 1e-6 级别

避坑指南：我曾经在融合时忘记把 BN 设成 eval 模式，结果 running_mean 和 running_var 还在更新，融合后的权重完全不对。另外，如果卷积层没有 bias，融合后一定要显式加上 bias，否则 BN 的 beta 没地方放。

至于 Conv+ReLU 的融合，在 PyTorch 里可以用 torch.nn.intrinsic 模块：

import torch.nn.intrinsic as nni

# 创建一个 ConvReLU2d 融合模块
conv_relu = nni.ConvReLU2d(
    nn.Conv2d(3, 16, 3, padding=1),
    nn.ReLU(inplace=True)
)

# 或者从已有的 Conv+ReLU 序列转换
conv = nn.Conv2d(3, 16, 3, padding=1)
relu = nn.ReLU(inplace=True)
fused = nni.ConvReLU2d(conv, relu)

这个融合模块在 forward 时，会在卷积计算完成后直接调用 ReLU，省掉了中间 tensor 的分配和拷贝。

总结一下

今天讲了两个最基础的融合：

  • Conv+BN 融合：把 BN 的线性变换吸收到卷积的权重和偏置里。数学上就是乘法和加法合并。
  • Conv+ReLU 融合：在卷积计算后直接做 in-place 的 ReLU，减少内存访问。

这两种融合在几乎所有推理框架中都是标配。但理解原理很重要，因为后面我们会讲到更复杂的融合，比如 Conv+BN+ReLU 三合一、残差分支融合等。那些都是在今天的基础上扩展的。

下节课我们会继续讲算子融合的下半部分，包括多分支融合和更高级的图优化技巧。到时候我会分享一个我在实际项目中踩过的坑——融合后精度反而下降了，你猜猜是什么原因？

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