读 ncnn 源码（XX）：`fuse_convolutiondepthwise_batchnorm`——为深度可分离卷积“减负”

在第十七篇中，我们深入分析了标准 Convolution 与 BatchNorm 的融合机制。深度可分离卷积（Depthwise Separable Convolution）作为现代轻量级网络（如 MobileNet, EfficientNet）的基础构件，其 ConvolutionDepthWise 部分后通常也紧跟着 BatchNorm 层。因此，针对性地优化 ConvolutionDepthWise -> BatchNorm 这一常见模式，对于提升这些模型的推理性能至关重要。

本篇，我们将剖析 fuse_convolutiondepthwise_batchnorm 的源码，理解 ncnn 如何将针对标准卷积的融合原理，平滑地应用于深度可分离卷积的 Depthwise 部分。

TL;DR

目标: 将 ConvolutionDepthWise 层后紧随的 BatchNorm 层进行融合，消除 BatchNorm 在推理时的计算开销。
模式匹配: 遍历网络 layers，查找 ConvolutionDepthWise -> BatchNorm 的直接连接模式。
数学原理: 与标准卷积融合 BN 完全一致。因为 ConvolutionDepthWise 的计算也是线性的（虽然是逐通道独立进行），BatchNorm 同样是线性的。融合公式依然是将 BN 的仿射变换参数（由 mean, var, slope, bias 导出）合并到 ConvolutionDepthWise 的权重和偏置中。
- 令 BN 参数导出系数 b = slope / sqrt(var + eps) 和 a = bias - slope * mean / sqrt(var + eps)。
- 新权重 $W'_{dw} = W_{dw} \cdot b$ （逐通道相乘）。
- 新偏置 $b'_{dw} = b_{dw} \cdot b + a$ （逐通道计算）。
代码实现: 几乎与 fuse_convolution_batchnorm 完全相同。
- 计算中间系数 a 和 b。
- 遍历 ConvolutionDepthWise 的每个通道 i (等价于 group 或 num_output)。
- 将该通道对应的权重 (weight_per_outch 个元素) 乘以 b[i]。
- 更新该通道的偏置 bias[i] = bias[i] * b[i] + a[i] （如果原层无偏置，则先创建）。
图结构修改: 将 ConvolutionDepthWise 的 top 指向原 BatchNorm 的 top，更新 blob 的 producer，并将 BatchNorm 标记为 "ncnnfused"。
效果: 消除了 BatchNorm 层的计算和内存访问，对于大量使用 Depthwise Convolution 的轻量级网络，性能提升显著。

1. 融合动机：通用线性合并原理的应用

特性	标准卷积 (Standard Conv)	深度卷积 (Depthwise Conv)
卷积核大小	$K \times K \times C_{\text{in}}$	$K \times K \times 1$
通道数	$C_{\text{out}}$ 个这样的卷积核	$C_{\text{in}}$ 个这样的卷积核
计算方式	跨通道（对所有输入通道加权求和，得到一个输出通道）	通道独立（每个核只负责一个输入通道，得到一个输出通道）
核是否一致	$C_{\text{out}}$ 个核相互独立，不一致。	$C_{\text{in}}$ 个核相互独立，不一致。
与 BN 融合	BN 参数 $K_i$ 和 $\mathbf{W}_i$ 作用于所有输入通道。	BN 参数 $K_i$ 和 $\mathbf{W}_{dw, i}$ 只作用于第 $i$ 个输入通道。

数学原理

1. 原始操作的数学表达式

我们关注网络中第 $i$ 个通道的计算，因为在深度卷积中，每个通道是独立处理的。

1.1. 深度卷积（ConvolutionDepthWise）

对于第 $i$ 个输入通道 $x_i$ ，深度卷积的输出 $y_i$ 为：

$y_i = \mathbf{W}_{dw, i} * \mathbf{x}_i + b_{dw, i}$

其中：

$\mathbf{W}_{dw, i}$ 是第 $i$ 个通道对应的卷积核（深度卷积的核是 $K \times K \times 1$ ）。
$b_{dw, i}$ 是第 $i$ 个通道的可选偏置。

1.2. 批量归一化（BatchNorm）

后续的 $\text{BatchNorm}$ 层对 $y_i$ 进行归一化和仿射变换，得到最终输出 $z_i$ ：

$z_i = \gamma_i \frac{y_i - \mu_i}{\sqrt{\sigma^2_i + \epsilon}} + \beta_i$

其中：

$\mu_i$ 和 $\sigma^2_i$ 是 $\text{BN}$ 层在训练过程中（或预先计算）确定的第 $i$ 个通道的均值和方差（用于推理）。
$\gamma_i$ 是第 $i$ 个通道的缩放因子（scale，也称 $\text{slope}$ ）。
$\beta_i$ 是第 $i$ 个通道的平移因子（shift，也称 $\text{bias}$ ）。
$\epsilon$ 是防止除以零的微小常数。

2. 融合的数学推导

我们的目标是将 $y_i$ 的表达式代入 $z_i$ 的表达式，然后将结果重写为新的深度卷积形式：

$z_i = \mathbf{W}'_{dw, i} * \mathbf{x}_i + b'_{dw, i}$

2.1. 代入 $y_i$

将 $\text{ConvDepthWise}$ 的输出 $y_i$ 代入 $\text{BatchNorm}$ 的公式中：

$z_i = \gamma_i \frac{(\mathbf{W}_{dw, i} * \mathbf{x}_i + b_{dw, i}) - \mu_i}{\sqrt{\sigma^2_i + \epsilon}} + \beta_i$

2.2. 引入常数因子 $K_i$

首先定义一个通道特定的归一化和缩放因子 $K_i$ ：

$K_i = \frac{\gamma_i}{\sqrt{\sigma^2_i + \epsilon}}$

$K_i$ 是一个常数，它完全取决于 $\text{BN}$ 层的参数。

2.3. 展开并重组

将 $K_i$ 代入 $z_i$ 的表达式并展开：

$z_i = K_i \cdot [(\mathbf{W}_{dw, i} * \mathbf{x}_i + b_{dw, i}) - \mu_i] + \beta_i$

$z_i = K_i \cdot (\mathbf{W}_{dw, i} * \mathbf{x}_i) + K_i \cdot b_{dw, i} - K_i \cdot \mu_i + \beta_i$

由于 $K_i$ 是一个常数，且卷积操作（ $*$ ）满足结合律和分配律（对于常数乘法），我们可以将 $K_i$ 移入卷积核中：

$z_i = (\mathbf{W}_{dw, i} \cdot K_i) * \mathbf{x}_i + \underbrace{(b_{dw, i} \cdot K_i - \mu_i \cdot K_i + \beta_i)}_{\text{新的偏置 } b'_{dw, i}}$

3. 融合后的参数

通过比较 $z_i = \mathbf{W}'_{dw, i} * \mathbf{x}_i + b'_{dw, i}$ ，我们得到融合后的新参数：

3.1. 融合后的权重 $\mathbf{W}'_{dw, i}$

新的深度卷积核 $\mathbf{W}'_{dw, i}$ 是原始卷积核与归一化缩放因子的逐元素乘积：

$\mathbf{W}'_{dw, i} = \mathbf{W}_{dw, i} \cdot K_i = \mathbf{W}_{dw, i} \cdot \frac{\gamma_i}{\sqrt{\sigma^2_i + \epsilon}}$

3.2. 融合后的偏置 $b'_{dw, i}$

新的偏置 $b'_{dw, i}$ 包含了原始偏置、归一化、缩放和平移的影响：

$b'_{dw, i} = (b_{dw, i} - \mu_i) \cdot K_i + \beta_i$

即：

$b'_{dw, i} = (b_{dw, i} - \mu_i) \cdot \frac{\gamma_i}{\sqrt{\sigma^2_i + \epsilon}} + \beta_i$

2. 代码实现：与标准卷积融合的高度相似性

fuse_convolutiondepthwise_batchnorm 的实现代码与 fuse_convolution_batchnorm 几乎如出一辙，这印证了其背后共通的数学原理。

int NetOptimize::fuse_convolutiondepthwise_batchnorm()
{
    const size_t layer_count = layers.size();
    for (size_t i = 0; i < layer_count; i++) // 遍历查找 ConvolutionDepthWise
    {
        if (layers[i]->type != "ConvolutionDepthWise") continue;
        int top_blob_index = layers[i]->tops[0];

        size_t j = i + 1; // 查找后续的 BatchNorm
        for (; j < layer_count; j++)
        {
            if (layers[j]->type != "BatchNorm") continue;
            if (layers[j]->bottoms.size() != 1) continue;
            if (layers[j]->bottoms[0] == top_blob_index) break; // 确认连接
        }
        if (j == layer_count) continue; // 未找到

        ncnn::ConvolutionDepthWise* convolutiondepthwise = (ncnn::ConvolutionDepthWise*)layers[i];
        ncnn::BatchNorm* batchnorm = (ncnn::BatchNorm*)layers[j];

        fprintf(stderr, "fuse_convolutiondepthwise_batchnorm %s %s\n", convolutiondepthwise->name.c_str(), batchnorm->name.c_str());

        // --- 参数变换核心 ---
        {
            int channels = batchnorm->channels; // = convolutiondepthwise->num_output
            float eps = batchnorm->eps;

            // 1. 计算中间系数 a 和 b (与 Conv+BN 融合完全相同)
            std::vector<float> a(channels);
            std::vector<float> b(channels);
            for (int ch = 0; ch < channels; ch++)
            {
                float sqrt_var = static_cast<float>(sqrt(batchnorm->var_data[ch] + eps));
                a[ch] = batchnorm->bias_data[ch] - batchnorm->slope_data[ch] * batchnorm->mean_data[ch] / sqrt_var;
                b[ch] = batchnorm->slope_data[ch] / sqrt_var;
            }

            // 2. 如果 DWConv 没有偏置，则创建并初始化为 0
            if (convolutiondepthwise->bias_term == 0)
            {
                convolutiondepthwise->bias_term = 1;
                convolutiondepthwise->bias_data = ncnn::Mat(channels);
                convolutiondepthwise->bias_data.fill(0.f);
            }

            // 3. 计算每个通道 (group) 的权重数量
            // 对于 DWConv, weight_data_size = channels * kernel_w * kernel_h
            // 所以 weight_per_outch = kernel_w * kernel_h
            const int weight_per_outch = convolutiondepthwise->weight_data_size / channels;

            float* weight = convolutiondepthwise->weight_data; // 指向 DWConv 权重
            float* bias = convolutiondepthwise->bias_data;   // 指向 DWConv 偏置

            // 4. 遍历每个通道 (group)，更新权重和偏置
            for (int ch = 0; ch < channels; ch++)
            {
                // 获取指向当前通道 ch 的权重的指针
                float* conv_weight_outch = weight + weight_per_outch * ch;
                // 将该通道的所有权重乘以 b[ch]
                for (int k = 0; k < weight_per_outch; k++)
                {
                    conv_weight_outch[k] *= b[ch]; // W'_dw = W_dw * b
                }
                // 更新偏置 b'_dw = b_dw * b + a
                bias[ch] = bias[ch] * b[ch] + a[ch];
            }
        } // --- 参数变换结束 ---

        // --- 图结构修改 (与 Conv+BN 融合完全相同) ---
        int top_blob_index_final = batchnorm->tops[0];
        convolutiondepthwise->tops[0] = top_blob_index_final;
        blobs[top_blob_index_final].producer = i;
        batchnorm->type = "ncnnfused";
        // --- 图结构修改结束 ---
    }
    return 0;
}

关键点:

模式匹配: 与标准卷积类似，查找 ConvolutionDepthWise -> BatchNorm 的直接连接。
参数计算: 计算 a, b 系数的方式与标准卷积完全一样。
权重更新: 由于 Depthwise 卷积的权重是 [channels, 1, H, W]（ncnn 内部可能是 [channels, H*W]），weight_per_outch 就是 H*W。代码正确地将每个通道 ch 对应的 H*W 个权重乘以了该通道的 b[ch]。
偏置更新: 与标准卷积一样，将偏置 bias[ch] 更新为 bias[ch] * b[ch] + a[ch]。
图修改: 与标准卷积一样，重定向连接并标记 BatchNorm 为无效。

3. 意义：轻量级网络的关键优化

深度可分离卷积将标准卷积分解为 Depthwise 和 Pointwise (1x1) 两步，大大减少了计算量和参数量，是构建轻量级网络的基石。然而，这两步之后通常都会接 BatchNorm 层。因此，能够融合 ConvolutionDepthWise + BatchNorm 和 Convolution (1x1) + BatchNorm 对于优化这类网络的性能至关重要。

fuse_convolutiondepthwise_batchnorm Pass 正是针对前者进行了优化，它与 fuse_convolution_batchnorm (可以处理 Pointwise 卷积) 协同工作，能够显著去除 MobileNet 等结构中大量的冗余 BatchNorm 计算。

4. 结语

fuse_convolutiondepthwise_batchnorm 体现了 ncnn 图优化策略的通用性和一致性。它将标准卷积与 BatchNorm 融合的成熟原理，直接应用于深度可分离卷积的 Depthwise 部分，代码实现高度复用，逻辑清晰。这一优化对于提升基于深度可分离卷积的轻量级模型在端侧设备上的推理速度具有重要意义。通过这一系列针对不同卷积类型（标准、深度、反卷积）的融合 Pass，ncnn 能够系统性地消除推理阶段不必要的 BatchNorm 计算，展现了其作为高性能推理框架的工程实力。

该封面图片由Nicolas IZERN在Pixabay上发布