读 ncnn 源码（XVIII）：`fuse_convolution_mul`——融合逐通道乘法，优化线性计算链

在上一篇中，我们分析了 fuse_convolution_batchnorm 如何通过代数等价变换，将 BatchNorm 层融入前驱的 Convolution 层。这一优化思路——合并连续的线性计算——具有普遍性。本篇，我们将探讨 ncnn 中另一个重要的融合 Pass：fuse_convolution_mul，它针对的是 **Convolution 层后紧跟一个特定模式的逐通道乘法（通过 BinaryOp 实现）**的场景。

这种 Conv -> Mul(per-channel) 的结构在某些网络（例如 MobileNetV1 中的 Depthwise + Pointwise 结构有时会伴随缩放）或模型转换过程中可能出现。将其融合，可以进一步减少计算冗余，提升推理效率。

TL;DR

目标: 将 Convolution 层后接一个扮演逐通道乘法 (Per-Channel Scaling) 角色的 BinaryOp 层（操作类型为 Mul，且第二个输入来自 MemoryData）进行融合。
模式匹配: 查找 Convolution -> BinaryOp 结构，并附加严格条件：
- BinaryOp 的类型必须是乘法 (op_type == 2)。
- BinaryOp 不能是与标量相乘 (!with_scalar)。
- BinaryOp 的第二个输入 bottoms[1] 必须来自一个 MemoryData 层。
- 该 MemoryData 层的数据维度必须符合逐通道乘法的广播要求（例如，w == num_output, h == 0, c == 0）。
数学原理:
- Conv: $y = W * x + b_{conv}$
- BinaryOp(Mul): $z = y \odot S$ (其中 $S$ 是来自 MemoryData 的、维度为 [num_output] 的缩放向量，⊙ 表示逐元素乘法，S 会被广播到 $y$ 的空间维度)
- 融合后: $z = W' * x + b'_{conv}$
- 推导得: $W'_{o,i,h,w} = W_{o,i,h,w} \cdot S_o$ (卷积核的每个输出通道 o 的所有权重乘以对应的缩放系数 $S_o$ )
- $b'_{conv, o} = b_{conv, o} \cdot S_o$ (卷积偏置的每个元素也乘以对应的缩放系数 $S_o$ )
代码实现:
- 遍历 Convolution 层的每个输出通道 i (0 to channels-1)。
- 将该通道对应的所有权重 (weight_per_outch 个元素) 乘以 memorydata->data[i]。
- 如果 Convolution 层有偏置 (bias_term == 1)，则将其偏置 bias[i] 也乘以 memorydata->data[i]。
图结构修改: 将 Convolution 层的 top 指向原 BinaryOp 层的 top，更新 blob 的 producer，并将 BinaryOp 层标记为 "ncnnfused"。
效果: 消除了 BinaryOp 层引入的逐元素乘法计算和额外的内存读写，将缩放操作无损地合并到卷积计算中。

1. 融合动机：消除冗余的逐通道缩放

Convolution 层本身是一个线性变换。如果其后紧跟一个 BinaryOp 层执行逐通道的乘法（即用一个向量乘以特征图的每个通道），那么这两个连续的线性操作在代数上是可以合并的。

BinaryOp 层虽然灵活，但在执行逐元素乘法时，需要遍历整个特征图，带来显著的内存访问开销。如果这个乘法仅仅是用一个固定的向量（来自 MemoryData）进行缩放，那么将这个缩放因子提前作用于 Convolution 层的权重和偏置，就可以在推理时完全省去 BinaryOp 的计算，同时得到完全相同的结果。

2. 模式匹配：精确锁定目标结构

fuse_convolution_mul 的模式匹配比 fuse_convolution_batchnorm 更为严格，因为它必须确保 BinaryOp 确实是在执行“逐通道乘法”这一特定任务。

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

        size_t j = i + 1; // 从下一层开始查找 BinaryOp
        for (; j < layer_count; j++)
        {
            if (layers[j]->type != "BinaryOp") continue;
            if (layers[j]->bottoms.size() != 2) continue; // BinaryOp 必须有两个输入
            if (layers[j]->bottoms[0] == top_blob_index) break; // 确认第一个输入来自 Conv
        }
        if (j == layer_count) continue; // 未找到

        ncnn::Convolution* convolution = (ncnn::Convolution*)layers[i];
        ncnn::BinaryOp* binaryop = (ncnn::BinaryOp*)layers[j];

        // 条件 1: 必须是乘法操作 (op_type 2 代表 Mul)
        // 条件 2: 不能是与标量相乘 (with_scalar 表示第二个操作数是常数)
        if (binaryop->op_type != 2 || binaryop->with_scalar) continue;

        // 条件 3: 查找 BinaryOp 的第二个输入的生产者
        size_t k = 0;
        for (; k < j; k++) // 必须在 BinaryOp 之前
        {
            if (layers[k]->type != "MemoryData") continue; // 必须是 MemoryData
            if (layers[k]->tops[0] == binaryop->bottoms[1]) break; // 确认连接关系
        }
        if (k == j) continue; // 未找到 MemoryData

        ncnn::MemoryData* memorydata = (ncnn::MemoryData*)layers[k];

        // 条件 4: 校验 MemoryData 的形状是否符合逐通道缩放 (向量)
        int channels = convolution->num_output;
        if (memorydata->w != channels || memorydata->h != 0 || memorydata->c != 0)
        {
            // 如果形状不匹配 (e.g., 是一个 Feature Map 而不是 Vector)，则不能融合
            continue;
        }

        // 所有条件满足，可以执行融合
        fprintf(stderr, "fuse_convolution_mul %s %s\n", convolution->name.c_str(), binaryop->name.c_str());
        // ... 执行参数变换和图修改 ...
    }
    return 0;
}

这个匹配过程确保了只有当 BinaryOp 确实扮演着“使用一个固定向量 S 对来自 Convolution 的特征图 y 进行逐通道相乘 y * S”的角色时，融合才会发生。

3. 数学原理与代码实现：吸收缩放因子

严格数学推理：吸收缩放因子

我们的目标是证明：

$(W * x + b_{conv}) \odot S = (W' * x) + b'_{conv}$

其中 $W'$ 和 $b'_{conv}$ 是融合了 $S$ 的新参数。

1. 定义与符号

$x$ : 输入张量，维度为 $(N, C_{in}, H_{in}, W_{in})$
$W$ : 卷积核权重，维度为 $(C_{out}, C_{in}, K_h, K_w)$
$b_{conv}$ : 卷积偏置，维度为 $(C_{out})$
$S$ : 缩放因子向量，维度为 $(C_{out})$
$*$ : 卷积操作
$\odot$ : 逐元素乘法（Hadamard 积）

2. 卷积的定义

首先，我们写出卷积层（带偏置）的输出 $Y = W * x + b_{conv}$ 中任意一个元素 $Y_{n, o, i, j}$ 的计算公式。

（ $n$ 为 batch 索引， $o$ 为输出通道索引， $i, j$ 为空间位置索引）

$Y_{n, o, i, j} = \left( \sum_{c=0}^{C_{in}-1} \sum_{k_h=0}^{K_h-1} \sum_{k_w=0}^{K_w-1} W_{o, c, k_h, k_w} \cdot x_{n, c, i', j'} \right) + b_{conv, o}$

（为简洁起见， $i', j'$ 代表 $x$ 中与 $W$ 的 $(k_h, k_w)$ 对应的输入位置，已考虑了步长(stride)和填充(padding)）

3. 应用缩放因子 $S$

接下来，我们计算最终输出 $z = Y \odot S$ 。

缩放因子 $S$ 的维度是 $(C_{out})$ 。当它与 $Y$ （维度为 $(N, C_{out}, H_{out}, W_{out})$ ）进行逐元素乘法 $\odot$ 时， $S$ 会被广播 (Broadcast) 到 $N, H_{out}, W_{out}$ 维度。

这意味着， $z$ 中的每一个元素 $z_{n, o, i, j}$ 的计算如下：

$z_{n, o, i, j} = Y_{n, o, i, j} \cdot S_o$

这是关键步骤： $Y$ 的第 $o$ 个输出通道上的所有元素，都乘以 $S$ 向量中的第 $o$ 个元素 $S_o$ 。

4. 代入与变换

现在，我们将第 (2) 步中的 $Y_{n, o, i, j}$ 表达式代入第 (3) 步：

$z_{n, o, i, j} = \left( \left( \sum_{c=0}^{C_{in}-1} \sum_{k_h=0}^{K_h-1} \sum_{k_w=0}^{K_w-1} W_{o, c, k_h, k_w} \cdot x_{n, c, i', j'} \right) + b_{conv, o} \right) \cdot S_o$

利用乘法对加法的分配律 (Distributive Property)，我们将 $S_o$ 乘入括号内：

$z_{n, o, i, j} = \left( \sum_{c=0}^{C_{in}-1} \sum_{k_h=0}^{K_h-1} \sum_{k_w=0}^{K_w-1} W_{o, c, k_h, k_w} \cdot x_{n, c, i', j'} \right) \cdot S_o + (b_{conv, o} \cdot S_o)$

由于 $S_o$ 是一个标量（对于 $c, k_h, k_w$ 的求和循环来说），我们可以将其移入求和符号内部（乘法结合律）：

$z_{n, o, i, j} = \left( \sum_{c=0}^{C_{in}-1} \sum_{k_h=0}^{K_h-1} \sum_{k_w=0}^{K_w-1} (W_{o, c, k_h, k_w} \cdot S_o) \cdot x_{n, c, i', j'} \right) + (b_{conv, o} \cdot S_o)$

5. 定义新参数 $W'$ 和 $b'_{conv}$

观察上式，我们可以定义一组新的权重 $W'$ 和偏置 $b'_{conv}$ ：

新权重 $W'$ ：

$W'_{o, c, k_h, k_w} = W_{o, c, k_h, k_w} \cdot S_o$

（即，原权重 $W$ 的第 $o$ 个输出通道对应的所有滤波器，都乘以 $S_o$ ）
新偏置 $b'_{conv}$ ：

$b'_{conv, o} = b_{conv, o} \cdot S_o$

（即，原偏置 $b_{conv}$ 的第 $o$ 个元素乘以 $S_o$ ）

6. 结论

将 $W'$ 和 $b'_{conv}$ 的定义代回到第 (4) 步的方程中：

$z_{n, o, i, j} = \left( \sum_{c=0}^{C_{in}-1} \sum_{k_h=0}^{K_h-1} \sum_{k_w=0}^{K_w-1} W'_{o, c, k_h, k_w} \cdot x_{n, c, i', j'} \right) + b'_{conv, o}$

这个 $z_{n, o, i, j}$ 的表达式，严格等于使用新权重 $W'$ 和新偏置 $b'_{conv}$ 对 $x$ 进行卷积操作的定义。

因此，我们严格证明了：

$(W * x + b_{conv}) \odot S = (W' * x) + b'_{conv}$

其中：

$W'_{o,i,h,w} = W_{o,i,h,w} \cdot S_o$

$b'_{conv, o} = b_{conv, o} \cdot S_o$

b) 代码实现:

{
    const int weight_per_outch = convolution->weight_data_size / channels;
    float* weight = convolution->weight_data; // 指向卷积权重数据
    float* bias = convolution->bias_data;   // 指向卷积偏置数据 (可能为 NULL)

    for (int i = 0; i < channels; i++) // 遍历每个输出通道
    {
        // 获取指向当前输出通道 i 的所有权重的指针
        float* conv_weight_outch = weight + weight_per_outch * i;
        // 将该通道的所有权重乘以对应的缩放系数 S_i
        for (int j = 0; j < weight_per_outch; j++)
        {
            conv_weight_outch[j] *= memorydata->data[i]; // W' = W * S_i
        }

        // 如果存在偏置，也将偏置乘以对应的缩放系数 S_i
        if (bias)
        {
            bias[i] = bias[i] * memorydata->data[i]; // b' = b * S_i
        }
    }
}

代码直接修改了 convolution->weight_data 和 convolution->bias_data 的内容，将缩放操作的效果“烘焙”进了卷积层的参数中。

4. 图结构修改

与 fuse_convolution_batchnorm 类似，参数修改完成后，需要更新计算图的拓扑结构：

int top_blob_index_final = binaryop->tops[0]; // 获取 BinaryOp 原来的输出 Blob
convolution->tops[0] = top_blob_index_final; // 将 Conv 的输出直接连接到最终 Blob
blobs[top_blob_index_final].producer = i;    // 更新 Blob 的生产者为 Conv 层
binaryop->type = "ncnnfused";               // 标记 BinaryOp 层为无效

这确保了推理引擎在执行时会跳过已被融合的 BinaryOp 层。

5. 意义与限制

意义:
- 减少计算量: 消除了 BinaryOp 层对整个特征图进行逐元素乘法的计算。
- 减少内存访问: 避免了读取 MemoryData、读取 Conv 输出、写入 BinaryOp 输出的内存操作。
- 优化线性链: 对于 Conv -> BN -> Scale -> Mul(per-channel) 这样的长线性链，通过多次融合 Pass，可以将其最终合并为一个单一的 Convolution 层，极大提升效率。
限制:
- 仅限 Mul: 此 Pass 只处理乘法类型的 BinaryOp。对于加法 (Add) 等其他类型，有专门的 fuse_convolution_add 等 Pass 处理。
- 仅限 MemoryData: 只处理第二个操作数来自固定 MemoryData 的情况。如果 BinaryOp 的两个输入都是动态计算的特征图，则无法融合。
- 广播模式: 必须是符合逐通道缩放的广播模式（MemoryData 是向量）。

6. 结语

fuse_convolution_mul 是 ncnn 图优化工具箱中针对线性计算链优化的又一利器。它精确地识别并融合了“卷积 + 逐通道缩放”这一特定模式，通过简单的参数变换，在推理时消除了冗余的 BinaryOp 计算。与 fuse_convolution_batchnorm 等其他融合 Pass 协同工作，ncnnoptimize 能够有效地简化网络结构，减少计算和访存开销，为模型在资源受限的端侧设备上高效运行提供保障。

该图片由🌸♡💙♡🌸 Julita 🌸♡💙♡🌸在Pixabay上发布