3. 手动融合实现：基于PyTorch的Conv+BN融合代码实战

好，前面我们把原理讲清楚了。现在咱们直接上手写代码。

我个人习惯是，先跑通一个简单的demo，再把它塞进模型里。这样心里有底。

3.1 融合的核心公式

先回顾一下，Conv+BN融合到底在算什么？

说白了，就是把BN层的四个参数——gamma、beta、mean、var——合并到Conv的weight和bias里。

公式长这样：

# 融合后的卷积权重
w_fused = w_conv * (gamma / sqrt(var + eps))

# 融合后的卷积偏置
b_fused = (b_conv - mean) * (gamma / sqrt(var + eps)) + beta

嗯，这里要注意：如果原Conv层没有bias，那b_conv就是0。很多模型为了省计算，确实会去掉bias，因为BN层自带平移能力。

3.2 手写融合函数

我在项目中遇到过好几次这样的场景：模型训练完，要转成ONNX或者TensorRT，结果发现BN层拖慢了推理速度。于是我就写了个通用融合函数。

import torch
import torch.nn as nn

def fuse_conv_bn(conv, bn):
    """
    融合 Conv2d + BatchNorm2d
    conv: nn.Conv2d
    bn:   nn.BatchNorm2d
    返回: 融合后的 nn.Conv2d
    """
    # 获取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)

    # 处理卷积权重
    w_fused = conv.weight.data * scale.view(-1, 1, 1, 1)

    # 处理偏置
    if conv.bias is not None:
        b_fused = (conv.bias.data - mean) * scale + beta
    else:
        b_fused = -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  # 融合后一定有bias
    )

    fused_conv.weight.data = w_fused
    fused_conv.bias.data = b_fused

    return fused_conv

这段代码我用了不下二十次。你想想看，每次写完都省掉一个BN层的计算，推理速度能快5%-15%。

3.3 验证融合结果

代码写完了，怎么知道对不对？

我的做法是：随机生成输入，分别跑融合前和融合后的模型，对比输出差异。

# 测试代码
torch.manual_seed(42)

# 创建测试模型
conv = nn.Conv2d(3, 16, kernel_size=3, padding=1, bias=False)
bn = nn.BatchNorm2d(16)

# 设置为eval模式（重要！）
conv.eval()
bn.eval()

# 随机输入
x = torch.randn(4, 3, 32, 32)

# 融合前输出
with torch.no_grad():
    out_before = bn(conv(x))

# 融合
fused_conv = fuse_conv_bn(conv, bn)

# 融合后输出
with torch.no_grad():
    out_after = fused_conv(x)

# 对比差异
diff = (out_before - out_after).abs().max().item()
print(f"最大差异: {diff:.6f}")

# 一般差异在1e-5以下就算成功
assert diff < 1e-5, "融合结果偏差过大！"

跑完这段，如果diff在1e-6量级，恭喜你，融合成功了。那点微小差异来自浮点运算顺序不同，完全不影响精度。

3.4 替换模型中的算子

单个Conv+BN融合搞定了。接下来要替换整个模型。

我常用的方法是递归遍历模型的所有子模块，找到Conv+BN对，替换掉。

def fuse_model(model):
    """
    递归替换模型中的所有Conv+BN对
    """
    model.eval()
    
    # 收集需要替换的路径
    fuse_list = []
    
    for name, child in model.named_children():
        if isinstance(child, nn.Sequential):
            # 处理Sequential内部
            layers = list(child.children())
            for i in range(len(layers) - 1):
                if isinstance(layers[i], nn.Conv2d) and isinstance(layers[i+1], nn.BatchNorm2d):
                    fuse_list.append((name, i, i+1))
        else:
            # 递归处理
            fuse_model(child)
    
    # 执行替换
    for parent_name, conv_idx, bn_idx in fuse_list:
        parent = model._modules[parent_name]
        conv = parent[conv_idx]
        bn = parent[bn_idx]
        
        fused_conv = fuse_conv_bn(conv, bn)
        
        # 替换Conv，移除BN
        parent[conv_idx] = fused_conv
        parent[bn_idx] = nn.Identity()  # 用恒等层占位
    
    return model

这段代码有个坑：替换后模型的参数量没变，但计算量减少了。因为BN层的计算被省掉了。

3.5 实战效果对比

拿ResNet18做个测试，效果很明显：

精度完全不变，延迟和显存都有明显下降。这就是白嫖的性能提升。

嗯，代码部分就这些。下一节我们会讲怎么把融合后的模型导出成ONNX，以及如何用TensorRT做进一步优化。不过那是后话了。

先把你手头的模型跑一遍融合，看看能快多少？

想要更多资料，联系可可