读 ncnn 源码（XVI）：`ncnnoptimize`——神经网络图优化的“炼金术”

在之前的篇章中，我们已经深入了解了 ncnn 如何加载模型 (Net::load_param/model)、执行推理 (Extractor::extract) 以及内部精巧的设计模式 (Pimpl)。现在，我们将目光投向一个独立但至关重要的工具——ncnnoptimize。这个工具扮演着“炼金术士”的角色，它接收原始的 ncnn 模型文件，通过一系列图优化 (Graph Optimization) 操作，将其“精炼”成一个更小、更快、更高效的版本，而完全不改变模型的数学等价性。

本篇，我们将剖析 ncnnoptimize.cpp 的 main 函数流程，并以 fuse_batchnorm_scale 为例，初步探索 ncnn 图优化的基本原理和实现方式。

TL;DR

ncnnoptimize 的定位: 一个离线 (Offline) 模型优化工具。它读取 ncnn 的 .param 和 .bin 文件，执行一系列图优化 Pass，然后输出优化后的新 .param 和 .bin 文件。
核心机制: 基于图模式匹配与变换 (Graph Pattern Matching and Transformation)。它会遍历网络计算图，寻找可以被优化（如融合、消除、替换）的特定子图结构，然后修改网络结构和权重参数。
主要优化类型:
- 算子融合 (Operator Fusion): 将多个连续的算子合并成一个等效的算子，减少函数调用开销和内存访问。例如，fuse_convolution_batchnorm, fuse_convolution_activation。
- 算子消除 (Operator Elimination): 移除对计算结果没有影响的多余算子。例如，eliminate_dropout (推理时无效), eliminate_noop, eliminate_split (如果后续分支被优化掉)。
- 算子替换 (Operator Replacement): 将某些算子组合替换为更高效的等效算子。例如，replace_reduction_with_global_pooling。
NetOptimize 类: 继承自 ModelWriter（ModelWriter 又继承自 ncnn::Net），因此它本身就是一个 ncnn::Net 对象。这使得它可以方便地加载、访问和修改网络的 layers 和 blobs。图优化操作直接在加载到内存中的网络结构上进行。
fuse_batchnorm_scale 示例: 这是一个典型的算子融合优化。它查找计算图中 BatchNorm -> Scale 的连续模式，然后通过修改 BatchNorm 层的 slope 和 bias 参数，将 Scale 层的计算效果代数等效地融入 BatchNorm 中，最后将 Scale 层标记为无效 ("ncnnfused")，从而在推理时跳过该层。

1. `ncnnoptimize` 工具概览 (`main` 函数流程)

ncnnoptimize.cpp 的 main 函数展示了该工具的标准工作流程：

int main(int argc, char** argv)
{
    // 1. 解析命令行参数: 输入/输出 param/bin 文件路径, 优化标志 (flag), 可选的模型裁剪点
    // ... argument parsing ...

    // 2. 创建 NetOptimize 对象
    NetOptimize optimizer;

    // 3. 设置存储类型 (flag 控制是否使用 FP16 存储)
    optimizer.storage_type = (flag == 65536 || flag == 1) ? 1 : 0;

    // 4. 加载原始模型 (复用基类 ncnn::Net 的功能)
    optimizer.load_param(inparam);
    optimizer.load_model(inbin); // 支持 "null" 表示无权重文件

    // 5. 设置模型裁剪 (可选)
    optimizer.set_cutparam(cutstartname, cutendname);

    // 6. *** 执行一系列图优化 Pass ***
    //    a) 融合 (Fusion)
    optimizer.fuse_batchnorm_scale();
    optimizer.fuse_convolution_batchnorm();
    //    ... (大量融合操作) ...
    optimizer.fuse_binaryop_eltwise();
    //    b) 替换 (Replacement)
    optimizer.replace_reduction_with_global_pooling();
    optimizer.replace_prelu_with_leaky_relu();
    //    ...
    //    c) 消除 (Elimination)
    optimizer.eliminate_dropout();
    optimizer.eliminate_pooling1x1();
    //    ... (大量消除操作) ...
    optimizer.eliminate_orphaned_memorydata();

    // 7. 形状推断 (Shape Inference) - 优化后可能需要重新推断中间 blob 形状
    optimizer.shape_inference();

    // 8. 估算内存占用 (可选信息)
    optimizer.estimate_memory_footprint();

    // 9. 保存优化后的模型 (复用基类 ModelWriter 的功能)
    optimizer.save(outparam, outbin);

    return 0;
}

核心流程: 加载 -> 逐一应用优化 Pass -> 保存。每一个 optimizer.fuse_...(), optimizer.eliminate_...(), optimizer.replace_...() 调用都代表一次对内存中网络计算图的扫描和修改尝试。

2. `NetOptimize` 类的设计：继承带来的便利

NetOptimize 的类继承关系 NetOptimize : public ModelWriter : public ncnn::Net 是其设计的关键：

继承 ncnn::Net: 使得 NetOptimize 对象天生就拥有加载 (load_param, load_model)、访问 (layers, blobs) 网络结构和权重的所有能力。
继承 ModelWriter: 提供了保存 (save) 优化后模型到文件的能力，以及处理权重数据格式（FP32/FP16 存储，fwrite_weight_data）和模型裁剪 (set_cutparam) 的辅助功能。
自身实现: NetOptimize 类专注于实现各种具体的图优化算法（作为其成员函数）。

这种设计使得图优化操作可以直接在加载到内存中的 ncnn::Net 对象上进行，代码逻辑清晰且高效。

3. 图优化 Pass 示例：`fuse_batchnorm_scale`

我们以 fuse_batchnorm_scale 为例，深入理解图优化的基本模式：

int NetOptimize::fuse_batchnorm_scale()
{
    // 1. 遍历网络中的每一层
    const size_t layer_count = layers.size();
    for (size_t i = 0; i < layer_count; i++)
    {
        // 2. 模式匹配：寻找 BatchNorm 层
        if (layers[i]->type != "BatchNorm")
            continue;

        int top_blob_index = layers[i]->tops[0]; // 获取 BatchNorm 的输出 blob 索引

        // 3. 模式匹配：在 BatchNorm 之后寻找直接消费其输出的 Scale 层
        size_t j = i + 1;
        for (; j < layer_count; j++)
        {
            if (layers[j]->type != "Scale") continue;
            if (layers[j]->bottoms.size() != 1) continue; // Scale 必须只有一个输入
            if (layers[j]->bottoms[0] == top_blob_index) break; // 找到了！
        }

        if (j == layer_count) continue; // 没有找到匹配的 Scale 层

        // 4. 获取匹配到的层对象指针
        ncnn::BatchNorm* batchnorm = (ncnn::BatchNorm*)layers[i];
        ncnn::Scale* scale = (ncnn::Scale*)layers[j];

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

        // 5. 参数变换：修改 BatchNorm 层的参数以吸收 Scale 层的功能
        {
            // BatchNorm: y = (x - mean) / sqrt(var + eps) * slope + bias
            // Scale:     z = y * scale_weight + scale_bias (if bias_term)
            // Fused BN:  z = (x - mean) / sqrt(var + eps) * (slope * scale_weight) + (bias * scale_weight + scale_bias)
            // 令: new_slope = slope * scale_weight
            //     new_bias  = bias * scale_weight + scale_bias

            int channels = batchnorm->channels;
            float* slope = batchnorm->slope_data; // BatchNorm 的 slope 参数
            float* bias = batchnorm->bias_data;   // BatchNorm 的 bias 参数

            for (int q = 0; q < channels; q++)
            {
                // 更新 slope
                slope[q] = slope[q] * scale->scale_data[q];
                // 更新 bias (考虑 Scale 是否有 bias)
                if (scale->bias_term)
                    bias[q] = bias[q] * scale->scale_data[q] + scale->bias_data[q];
                else
                    bias[q] = bias[q] * scale->scale_data[q];
            }
        }

        // 6. 图结构修改：
        //    a) 将 BatchNorm 的输出直接连接到原 Scale 的输出 blob
        int top_blob_index_final = scale->tops[0];
        batchnorm->tops[0] = top_blob_index_final;
        //    b) 更新最终输出 blob 的生产者信息
        blobs[top_blob_index_final].producer = i; // 新的生产者是 BatchNorm 层
        //    c) 标记 Scale 层为无效/已融合 (ncnn 中常用修改 type 的方式)
        scale->type = "ncnnfused"; // 这个类型在推理时会被跳过
    }

    return 0;
}

核心步骤:

遍历与模式匹配: 找到 BatchNorm -> Scale 这样的相邻层结构。
参数等效变换: 基于数学原理，修改 BatchNorm 层的权重 (slope_data, bias_data)，使其计算结果等效于原始的 BatchNorm + Scale。
图结构修改: 断开 BatchNorm 和 Scale 的连接，将 BatchNorm 的输出直接连接到原 Scale 的下一层，并将 Scale 层标记为无效。

效果: 在推理时，Scale 层会被完全跳过，减少了一次层计算的开销（包括函数调用、内存读写等），同时保持了计算结果的数学等价性。

fuse_batchnorm_scale（融合 BatchNorm 和 Scale 层）是一种在模型**推理阶段（Inference）**进行的优化。

其核心原理是：将两个连续的、纯数学运算的层（BatchNorm 和 Scale）合并成一个单独的层，从而减少计算量和内存访问次数。

这种融合是数学上等价的，它通过预先计算一个新的缩放（Scale）和偏置（Bias）参数，来替代原有两个层的计算。

详细原理分解

要理解融合，我们首先要看这两个层在**推理（Inference）**时分别做了什么。

假设一个特征图（Feature Map）的某个通道的输入为 $x$ 。

1. BatchNorm 层的计算

在训练时，BatchNorm 会计算当前批次（batch）的均值和方差。但在推理时，它使用的是训练过程中累积的全局均值（running mean） $\mu$ 和全局方差（running variance） $\sigma^2$ 。这两个值是固定不变的常量。

BatchNorm 层的计算公式为：

$x_{norm} = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}}$

$x$ ：输入数据
$\mu$ ：全局均值（常量）
$\sigma^2$ ：全局方差（常量）
$\epsilon$ ：一个极小的数（常量），防止除以零
$x_{norm}$ ：归一化后的输出

注意： 在某些框架（如 Caffe）中，BatchNorm 层只做上述的归一化操作。

2. Scale 层的计算

在 Caffe 这类框架中，BatchNorm 层之后会紧跟一个 Scale 层。这个层的作用是进行缩放和平移，以恢复网络的表达能力（因为归一化会限制数据的分布）。

Scale 层有两个可学习的参数：缩放因子 $\gamma$ （gamma）和偏置（或平移）因子 $\beta$ （beta）。在推理时，这两个参数也是固定不变的常量。

Scale 层的计算公式为：

$y = \gamma \cdot x_{norm} + \beta$

$x_{norm}$ ：BatchNorm 层的输出
$\gamma$ ：缩放因子（常量）
$\beta$ ：偏置因子（常量）
$y$ ：最终的输出

3. 融合（Fuse）的数学原理

融合的目标是找到一组新的 $\gamma_{new}$ 和 $\beta_{new}$ ，使得 $y = \gamma_{new} \cdot x + \beta_{new}$ 这个单层计算与上述两层计算的结果完全等价。

我们将 BatchNorm 的公式代入 Scale 的公式中：

$y = \gamma \cdot \left( \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \right) + \beta$

现在，我们把这个公式展开并重新整理，使其变成 $A \cdot x + B$ 的形式：

$y = \left( \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}} \right) \cdot (x - \mu) + \beta$

$y = \left( \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}} \right) \cdot x - \left( \frac{\gamma \cdot \mu}{\sqrt{\sigma^2 + \epsilon}} \right) + \beta$

$y = \underbrace{\left( \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}} \right)}_{\text{新的 } \gamma_{new}} \cdot x + \underbrace{\left( \beta - \frac{\gamma \cdot \mu}{\sqrt{\sigma^2 + \epsilon}} \right)}_{\text{新的 } \beta_{new}}$

我们成功地将两个计算合并成了一个。

新的 $\gamma_{new}$ (新Scale):

$\gamma_{new} = \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}}$
新的 $\beta_{new}$ (新Bias):

$\beta_{new} = \beta - \frac{\gamma \cdot \mu}{\sqrt{\sigma^2 + \epsilon}}$

在模型加载时， $\mu, \sigma^2, \gamma, \beta, \epsilon$ 全都是已知的常量。因此，ncnn (或任何推理引擎) 可以在加载模型时一次性计算出 $\gamma_{new}$ 和 $\beta_{new}$ 。

然后，引擎会丢弃原来的 BatchNorm 层，并用这对新的 $(\gamma_{new}, \beta_{new})$ 参数替换掉原来 Scale 层的 $(\gamma, \beta)$ 参数。

为什么这么做？（优化的好处）

减少计算量：
- 融合前： 每次推理需要 1 次减法、1 次开方、1 次除法、1 次乘法、1 次加法。
- 融合后： 每次推理只需要 1 次乘法和 1 次加法。
- (注：开方和除法通常是比乘法和加法昂贵得多的计算操作)。
减少内存访问：
- 融合前： 需要 CPU/GPU -> 读取 $x$ -> 计算并写回 $x_{norm}$ -> 读取 $x_{norm}$ -> 计算并写回 $y$ 。
- 融合后： 只需要 CPU/GPU -> 读取 $x$ -> 计算并写回 $y$ 。
- 这减少了一次对中间结果 $x_{norm}$ 的读写，在 GPU 上这种内存带宽（Memory Bandwidth）的节省非常可观。

进一步的融合：`Conv + BatchNorm + Scale`

fuse_batchnorm_scale 通常只是第一步。在实际应用中（如 ncnn），最常见的融合是 Conv -> BatchNorm -> Scale 三合一。

原理是相同的，只是数学上更进一步：

Conv 层计算： $x_{conv} = W \cdot x_{input} + b$
BN + Scale 计算 (已融合)： $y = \gamma_{new} \cdot x_{conv} + \beta_{new}$
代入：

$y = \gamma_{new} \cdot (W \cdot x_{input} + b) + \beta_{new}$
展开：

$y = (\gamma_{new} \cdot W) \cdot x_{input} + (\gamma_{new} \cdot b + \beta_{new})$
最终融合：
- 新权重 $W_{fused}$ : $W_{fused} = \gamma_{new} \cdot W$
- 新偏置 $b_{fused}$ : $b_{fused} = \gamma_{new} \cdot b + \beta_{new}$

推理引擎可以直接计算出一套新的、等价的卷积权重 $W_{fused}$ 和偏置 $b_{fused}$ ，然后只保留这个 Conv 层，把 BatchNorm 和 Scale 层彻底丢弃。这就是为什么在 ncnn 这类推理引擎中，BatchNorm 层在 create_pipeline 阶段就消失了。

4. 图优化的意义与价值

ncnnoptimize 所执行的图优化，对于提升模型在端侧设备上的性能至关重要：

减少计算量: 融合操作可以将多个计算步骤合并，减少总的浮点运算次数（虽然 BN+Scale 例子主要减少的是访存和调用开销）。
减少内存访问: 融合可以消除中间结果 blob 的写入和读取，降低内存带宽压力。
减少模型加载时间与体积: 消除冗余层可以使 .param 文件变小，减少解析时间。
提升缓存利用率: 融合后的算子通常具有更好的数据局部性。

虽然这些优化技巧看起来“微小”，但在神经网络这种计算密集的场景下，积少成多，最终能带来显著的性能提升（例如 10%-30% 甚至更高）。

5. 结语

ncnnoptimize 工具是 ncnn 生态系统中至关重要的一环。它通过对加载到内存中的 ncnn::Net 对象执行一系列基于图模式匹配与变换的优化 Pass，实现了对原始神经网络计算图的“精炼”。以 fuse_batchnorm_scale 为代表的算子融合技术，清晰地展示了如何通过代数等价变换，在保持模型精度的前提下，减少计算冗余，提升推理效率。理解这些图优化策略，不仅有助于我们更好地利用 ncnn 工具，也为我们理解更高级的 AI 编译器（如 TVM, MLIR）的优化原理打下了基础。

该封面图片由Heike Tönnemann在Pixabay上发布