读 ncnn 源码（XXX）：replace_reduction_with_global_pooling——算子替换：识别低效模式

在 ncnnoptimize 的优化 Pass 中，我们已经见证了大量“算子融合”（Operator Fusion）操作，它们通过合并相邻的线性层（如 Conv+BN）来减少计算和访存开销。本篇，我们将探讨另一种同样重要的优化技术：算子替换 (Operator Replacement)。replace_reduction_with_global_pooling 函数就是其典型代表。

它的核心任务是识别出一种低效的、由多个通用层组合而成的计算模式，并将其替换为一个功能等价但性能高得多的专用层。

TL;DR

  1. 目标: 识别并替换一种用两次连续的 Reduction (Mean) 操作来实现“全局平均池化”(Global Average Pooling) 的低效模式。
  2. 模式匹配: 遍历 layers，查找 Reduction -> Reduction 的直接连接模式。
  3. 严格条件: 两个 Reduction 层都必须是 Mean 操作 (operation == 3)，非全局 (reduce_all == 0)，且系数为 1 (coeff == 1.f)。此外，它们必须各自只对一个轴进行操作 (axes.w == 1)，并且轴的组合（如 axis=2 或 3 后接 axis=2，这可能对应于 H 和 W 轴）符合 ncnn 对特定框架（如 ONNX/TensorFlow 转换而来）全局池化模式的识别规则。
  4. 替换动作:
    • 动态创建一个全新的 Pooling 层 (ncnn::create_layer_cpu("Pooling"))。
    • 将新 Pooling 层的参数设置为 pooling_type = 1 (Average) 和 global_pooling = 1。
    • 新的 Pooling 层在 layers 向量中替换掉第二个 Reduction 层的 位置 (layers[j] = pooling)。
  5. 图结构修改:
    • 将第一个 Reduction 层标记为无效 (reduction1->type = "ncnnfused")。
    • 将新 Pooling 层的 bottom（输入）直接连接到第一个 Reduction 层的 bottom，彻底绕过这个双层 Reduction 结构。
  6. 效果: 这种替换是巨大的性能提升。它将两个通用的、可能基于循环实现的 Reduction 层的计算和访存开销，替换为一个专用的、内部有高效 SIMD 实现的 GlobalAveragePooling 层的开销。

1. 融合之外的另一条路——算子替换

图优化的目标是等价地提升性能。除了将相邻层合并（融合），另一种强大的技术是识别出一个由多个通用层组成的、意图上等价于某个专用层的模式，然后执行替换。

GlobalAveragePooling (GAP) 是一个非常常见的操作，它将特征图的 H 和 W 维度上的所有像素值取平均，最终输出一个 [C] 维度的向量。然而，一些深度学习框架（如 TensorFlow 或 ONNX 转换器）在导出模型时，可能会将这个操作“降级”为两次连续的 ReduceMean 操作：

  1. ReduceMean(axis=H)：对 H 维度（高）取均值。
  2. ReduceMean(axis=W)：对上一步的结果再对 W 维度（宽）取均值。

这个 Reduction -> Reduction 的组合在数学上等价于 GAP，但在 ncnn 中执行效率低下：

  • 两次层调度: 需要两次 forward_layer 调用。
  • 两次内存读写: 第一个 Reduction 将结果写入一个中间 blob，第二个 Reduction 再将其读出。
  • 通用实现: Reduction 层是一个通用的n维规约操作，其实现可能不如专用的 Pooling 层针对二维空间的全局池化优化得那么极致。

replace_reduction_with_global_pooling 的使命，就是识别出这个低效模式，并将其“升级”为 ncnn 内部最高效的 GlobalAveragePooling 实现。

2. 模式匹配：寻找连续的空间均值

该函数的匹配逻辑非常严格，旨在精确捕获上述模式：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
int NetOptimize::replace_reduction_with_global_pooling()
{
    const size_t layer_count = layers.size();
    for (size_t i = 0; i < layer_count; i++)
    {
        // 1. 查找第一个 Reduction 层
        if (layers[i]->type != "Reduction") continue;

        ncnn::Reduction* reduction1 = (ncnn::Reduction*)layers[i];
        
        // 2. 检查是否为 "ReduceMean(axis=H or W)" 的特定形式
        //    operation 3 似乎对应 Mean
        if (reduction1->operation != 3 || reduction1->reduce_all != 0 || reduction1->coeff != 1.f)
            continue;
        if (reduction1->axes.w != 1) continue; // 必须是单轴
        
        // ncnn 内部的轴索引，2 和 3 可能对应某些框架的 H 和 W 轴
        const int* axes_ptr = reduction1->axes;
        if (axes_ptr[0] != 2 && axes_ptr[0] != 3)
            continue;

        // 3. 查找紧随其后的第二个 Reduction 层
        int top_blob_index = layers[i]->tops[0];
        size_t j = i + 1;
        for (; j < layer_count; j++)
        {
            if (layers[j]->type != "Reduction") continue;
            if (layers[j]->bottoms.size() != 1) continue;
            if (layers[j]->bottoms[0] == top_blob_index) break; // 确认连接
        }
        if (j == layer_count) continue; // 未找到

        ncnn::Reduction* reduction2 = (ncnn::Reduction*)layers[j];

        // 4. 检查是否为 "ReduceMean(axis=W or H)" 的特定形式
        if (reduction2->operation != 3 || reduction2->reduce_all != 0 || reduction2->coeff != 1.f)
            continue;
        if (reduction2->axes.w != 1) continue;
        
        // 检查是否为另一个空间轴 (e.g., axis=2)
        const int* axes2_ptr = reduction2->axes;
        if (axes2_ptr[0] != 2)
            continue;

        // 5. 模式匹配成功
        fprintf(stderr, "replace_reduction_with_global_pooling %s %s\n", reduction1->name.c_str(), reduction2->name.c_str());
        
        // ... 执行替换 ...
    }
    return 0;
}
  • 模式: Reduction(op=Mean, axis=2 or 3) -> Reduction(op=Mean, axis=2)。
  • 意图: 这里的轴索引 2 和 3 很可能对应于特定模型格式（如 ONNX [N, C, H, W] 的 axes=[2, 3]）在 ncnn Mat 内部（可能是 [W, H, C] 或 [W, H, D, C]）的轴映射。代码的意图是匹配在两个空间维度上的连续 ReduceMean。

3. 执行替换：Pooling 层的“偷天换日”

一旦模式匹配成功，优化器会执行一次“偷天换日”：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
// 1. 创建一个全新的 Pooling 层
ncnn::Pooling* pooling = (ncnn::Pooling*)ncnn::create_layer_cpu("Pooling");

// 2. 继承原 Layer[j] 的基本信息（名称、输入、输出）
pooling->type = "Pooling";
pooling->name = reduction2->name;
pooling->bottoms = reduction2->bottoms; // 暂时继承
pooling->tops = reduction2->tops;

// 3. 加载空参数（Pooling 层通常有自己的参数设置方式）
ncnn::ParamDict pd;
pooling->load_param(pd);

// 4. 关键配置：将其设置为“全局平均池化”
pooling->pooling_type = 1; // 1 = PoolMethod_AVE
pooling->global_pooling = 1; // 1 = true

// 5. 在网络层列表中，用新层替换旧层
layers[j] = pooling;
delete reduction2; // 销毁原 Reduction[j] 对象

核心: pooling->global_pooling = 1 是关键。当 Pooling 层的这个标志位被设为 1 时，它的 forward 函数会忽略所有 kernel_w, stride 等参数，转而调用一个专门的、高度优化的 pooling_global_... 函数，该函数直接对整个 W 和 H 维度进行求平均。

4. 计算图重连接 (Graph Rewiring)

替换完成后，还需要修正计算图的连接，让这个新层真正生效：

1
2
3
4
5
6
7
8
9
10
11
// 1. 获取 reduction1 的原始输入
int bottom_blob_index_final = reduction1->bottoms[0];

// 2. 将新 Pooling 层的输入，直接指向 reduction1 的输入
pooling->bottoms[0] = bottom_blob_index_final;

// 3. 更新原始输入的消费者信息，指向新 Pooling 层 (索引 j)
blobs[bottom_blob_index_final].consumer = j;

// 4. 标记 reduction1 为无效，推理时将被跳过
reduction1->type = "ncnnfused";

效果: 原本 Blob_A -> Reduction1 -> Blob_B -> Reduction2 -> Blob_C 的计算流，被重塑为 Blob_A -> GlobalPooling -> Blob_C。Reduction1 和 Blob_B 在推理时被完全跳过。

5. 结语

replace_reduction_with_global_pooling 是一个典型的算子替换优化 Pass。它展示了 ncnnoptimize 不仅能进行局部的算子融合，还能识别出由多个低效、通用算子组成的、具有特定语义模式（如全局池化）的子图，并将其替换为单一的、高性能的专用算子。这种“理解”模型意图并进行重构的能力，是图优化工具提升模型性能的关键所在，它允许 ncnn 用自己最高效的 SIMD 内核（如 pooling_global_avx）来替代转换器产生的低效实现。