读 ncnn 源码（XVII）：fuse_convolution_batchnorm——融合 BN，轻装前行

在上一篇《读 ncnn 源码（XVI）》中，我们初步探索了 ncnnoptimize 工具及其图优化（Graph Optimization）的基本原理，并以 fuse_batchnorm_scale 为例展示了算子融合（Operator Fusion）的威力。本篇，我们将深入分析一个更常见、影响也更广泛的融合优化：卷积层（Convolution）与批归一化层（BatchNorm）的融合，即 fuse_convolution_batchnorm。

几乎所有的现代卷积神经网络在训练阶段都会大量使用 BatchNorm 层来加速收敛、提升泛化能力。然而，在推理阶段，BatchNorm 的计算是完全线性的，可以被等效地合并到其前面的线性层（如卷积层）中。理解 fuse_convolution_batchnorm 的原理与实现，对于掌握 ncnn 乃至所有推理引擎的核心优化技术至关重要。

TL;DR

  1. 目标: 将推理阶段计算固定、但开销不小的 BatchNorm 层，通过数学等价变换，将其参数融入到前驱的 Convolution 层中，从而在推理时完全消除 BatchNorm 层的计算。
  2. 模式匹配: fuse_convolution_batchnorm 函数遍历网络 layers，寻找 Convolution -> BatchNorm 这样的直接连接模式（即 BatchNorm 的唯一 bottom 是 Convolution 的唯一 top）。
  3. 数学原理:
    • Conv: y=W∗x+bconvy = W * x + b_{conv}
    • BN: z=γy−μσ2+ϵ+βz = \gamma \frac{y - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta (其中 γ\gamma 是 slope, β\beta 是 bias, μ\mu 是 mean, σ2\sigma^2 是 var)
    • 融合后: z=W′∗x+bconv′z = W' * x + b'_{conv}
    • 推导得: W′=W⋅γσ2+ϵW' = W \cdot \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}}
    • bconv′=(bconv−μ)⋅γσ2+ϵ+βb'_{conv} = (b_{conv} - \mu) \cdot \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}} + \beta
  4. 代码实现:
    • 计算中间系数 a 和 b：b = gamma / sqrt(var + eps), a = beta - gamma * mean / sqrt(var + eps)。
    • 更新卷积权重: conv_weight_outch[j] *= b[i] (逐元素乘以 b)。
    • 更新卷积偏置: bias[i] = bias[i] * b[i] + a[i]。如果原卷积层没有偏置 (bias_term == 0)，则先创建并初始化为零偏置。
  5. 图结构修改: 将 Convolution 层的 top 指向原 BatchNorm 层的 top，更新 blob 的 producer 信息，并将 BatchNorm 层的 type 修改为 "ncnnfused" 以在推理时跳过。
  6. 效果: 显著减少推理时的计算量（消除了 BN 的逐元素运算）和内存访问（消除了 BN 输出 blob 的读写），是部署阶段非常关键的性能优化。

1. 为何要融合 Conv + BN？

BatchNorm 层在训练时通过计算批次统计量（均值、方差）进行归一化，但在推理时，使用的是训练阶段得到的全局（或滑动平均）统计量 μ\mu 和 σ2\sigma^2，以及学习到的缩放系数 γ\gamma (slope_data) 和偏移系数 β\beta (bias_data)。此时，BatchNorm 的计算公式变为纯粹的线性变换：

zi=γiyi−μiσi2+ϵ+βiz_i = \gamma_i \frac{y_i - \mu_i}{\sqrt{\sigma^2_i + \epsilon}} + \beta_i

其中 yiy_i 是输入特征图的第 ii 个通道，ziz_i 是对应的输出通道。ϵ\epsilon 是一个防止除零的小常数。

虽然这个计算本身不复杂，但它作用于整个特征图的每一个像素上。对于移动端 CPU 而言，这意味着大量的内存读写（读取 yiy_i，写入 ziz_i）和逐元素运算（减法、除法、乘法、加法）。这些操作，尤其是内存访问，带来的开销不容忽视。

而卷积层 y = W * x + b_{conv} 本身也是一个线性运算。将两个连续的线性运算合并成一个，是代数优化的基本思路。

2. 融合的数学推导与代码实现

fuse_convolution_batchnorm 的核心就是将上述两个线性公式合并。

2.1 数学推导

将卷积公式代入 BN 公式：

z=γ(W∗x+bconv)−μσ2+ϵ+βz = \gamma \frac{(W * x + b_{conv}) - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta

整理得：

z=(W∗x)⋅γσ2+ϵ+(bconv−μ)⋅γσ2+ϵ+βz = \left( W * x \right) \cdot \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}} + (b_{conv} - \mu) \cdot \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}} + \beta

令新的卷积权重为 W′W'，新的卷积偏置为 bconv′b'_{conv}，则融合后的形式为 z=W′∗x+bconv′z = W' * x + b'_{conv}。对比系数可得：

W′=W⋅γσ2+ϵW' = W \cdot \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}}

bconv′=(bconv−μ)⋅γσ2+ϵ+βb'_{conv} = (b_{conv} - \mu) \cdot \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}} + \beta

2.2 代码实现

ncnn 的代码正是严格按照这个推导来实现的：

a) 计算中间系数 a 和 b:

代码为了简化计算，先计算了两个中间数组 a 和 b，对应我们推导中的部分项：

1
2
3
4
5
6
7
8
9
10
// a = bias - slope * mean / sqrt(var + eps)  => 对应推导中的 β - γ*μ / sqrt(...)
// b = slope / sqrt(var + eps)               => 对应推导中的 γ / sqrt(...)
std::vector<float> a(channels);
std::vector<float> b(channels);
for (int i = 0; i < channels; i++)
{
    float sqrt_var = static_cast<float>(sqrt(batchnorm->var_data[i] + eps));
    a[i] = batchnorm->bias_data[i] - batchnorm->slope_data[i] * batchnorm->mean_data[i] / sqrt_var;
    b[i] = batchnorm->slope_data[i] / sqrt_var;
}

这里的 a[i] 和 b[i] 是针对每个通道 i 计算的。

b) 更新卷积权重 W′=W⋅bW' = W \cdot b:

1
2
3
4
5
6
7
8
9
10
11
const int weight_per_outch = convolution->weight_data_size / channels; // 计算每个输出通道对应的权重数量
float* weight = convolution->weight_data;
for (int i = 0; i < channels; i++) // 遍历每个输出通道
{
    float* conv_weight_outch = weight + weight_per_outch * i; // 获取该输出通道对应的所有权重
    for (int j = 0; j < weight_per_outch; j++) // 遍历该通道的所有权重
    {
        conv_weight_outch[j] *= b[i]; // 将权重乘以对应的 b[i]
    }
    // ... 更新偏置 ...
}

注意，BatchNorm 的参数是 per-channel 的，所以同一个输出通道 i 对应的所有卷积权重（无论它们连接哪个输入通道或处于哪个空间位置）都需要乘以同一个 b[i]。

c) 更新卷积偏置 bconv′=bconv⋅b+ab'_{conv} = b_{conv} \cdot b + a:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
// 如果原卷积层没有偏置，先创建并初始化为 0
if (convolution->bias_term == 0)
{
    convolution->bias_term = 1;
    convolution->bias_data = ncnn::Mat(channels);
    convolution->bias_data.fill(0.f);
}
float* bias = convolution->bias_data;
for (int i = 0; i < channels; i++)
{
    // ... 更新权重 ...
    // 更新偏置
    bias[i] = bias[i] * b[i] + a[i]; // bias_new = bias_old * b + a
}

这里的 bias[i] 对应 bconvb_{conv}，b[i] 对应 γσ2+ϵ\frac{\gamma}{\sqrt{\sigma^2 + \epsilon}}，a[i] 对应 β−γμσ2+ϵ\beta - \frac{\gamma \mu}{\sqrt{\sigma^2 + \epsilon}}。代码完美匹配了推导公式 bconv′=bconv⋅b+ab'_{conv} = b_{conv} \cdot b + a。

2.3 图结构修改

参数更新完毕后，还需要修改计算图的连接关系：

1
2
3
4
int top_blob_index_final = batchnorm->tops[0]; // 获取 BN 原来的输出 Blob
convolution->tops[0] = top_blob_index_final; // 将 Conv 的输出直接指向最终的 Blob
blobs[top_blob_index_final].producer = i;    // 更新 Blob 的生产者为 Conv 层
batchnorm->type = "ncnnfused";               // 标记 BN 层为无效

这样，在推理时，引擎执行完 Convolution 层后，会根据 convolution->tops[0] 直接跳转到后续层，而 "ncnnfused" 类型的 BatchNorm 层会被完全跳过。

3. 模式匹配与遍历

与 fuse_batchnorm_scale 类似，fuse_convolution_batchnorm 也采用双层循环进行模式匹配：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
const size_t layer_count = layers.size();
for (size_t i = 0; i < layer_count; i++) // 外层循环遍历所有层
{
    if (layers[i]->type != "Convolution") continue; // 找到一个 Conv 层

    int top_blob_index = layers[i]->tops[0];

    size_t j = i + 1; // 内层循环从 Conv 层的下一层开始查找
    for (; j < layer_count; j++)
    {
        if (layers[j]->type != "BatchNorm") continue; // 找到一个 BN 层
        if (layers[j]->bottoms.size() != 1) continue; // 确保 BN 只有一个输入
        if (layers[j]->bottoms[0] == top_blob_index) break; // 确认 BN 的输入是 Conv 的输出
    }

    if (j == layer_count) continue; // 未找到匹配

    // ... 执行融合 ...
}

这种简单的线性扫描匹配适用于常见的顺序模型结构。

4. 结语

fuse_convolution_batchnorm 是 ncnnoptimize 中一项基础且极其有效的图优化 Pass。它通过精确的数学推导，将 BatchNorm 的线性计算无损地合并到前驱的 Convolution 层中，直接消除了 BatchNorm 在推理阶段的计算和访存开销。这种基于代数等价性的算子融合技术，是所有深度学习推理框架提升性能的通用手段。理解其原理，有助于我们认识到神经网络模型在训练和推理阶段存在的巨大优化空间，以及图优化在部署流程中的核心价值。

该封面图片由Andrei Sla在Pixabay上发布