读 ncnn 源码（XLI）：shape_inference——图优化的“沙盘推演”与“状态重整”

ncnnoptimize 的工作流是一系列对计算图的破坏性修改（融合、消除、替换）。这些操作的副作用是，许多中间 Blob 的维度信息（shape）可能会变得陈旧或不明确。shape_inference Pass 的职责并非“优化”，而是一个至关重要的**“工具型 Pass”**。

它的核心任务是：通过一次模拟的“沙盘推演”（Dry Run），在不执行实际数值计算的前提下，强制 ncnn 的推理引擎（Extractor）完整地跑一遍优化后的计算图，从而“推演”并“记录”下每一个 Blob 的确切形状信息。

TL;DR

  1. 目标: 在所有图优化 Pass 执行完毕后，为网络中的每一个 Blob 重新计算并填充其 shape 属性（dims, w, h, c, elempack），并将其回填到每个 Layer 的 bottom_shapes 和 top_shapes 成员中。
  2. 为何必须:
    • 状态更新: 优化 Pass（如 fuse_...）会改变 Layer 的参数（如 activation_type），shape_inference 中的 recreate_pipeline 步骤会强制 Layer 根据新参数重新生成其内部状态（如打包后的权重 weight_data_tm）。
    • 数据支撑: 后续的优化 Pass（如 estimate_memory_footprint）或最终 save 保存 .param 文件时，需要依赖这些精确的 Blob 形状信息。
  3. 核心机制 (Dry Run):
    • 重置管线: 遍历所有 Layer，调用 destroy_pipeline 和 create_pipeline，强制它们根据优化后的新参数重新初始化。
    • 创建提取器: create_extractor() 并设置为 light_mode。
    • 喂入“假”数据: 找到所有 Input 层，根据其必须已定义的 w, h, c 形状，创建“假”的（未初始化）Mat，并通过 ex.input() 注入。
    • 触发全图计算: 遍历所有 Layer，并对每个 Layer 的所有 top_blob 调用 ex.extract(top_blob_index, m)。
    • 惰性求值的妙用: ex.extract() 会触发 ncnn 的惰性求值机制，递归地计算所有依赖项。通过对 所有 top_blob 都调用一次 extract，ncnn 被迫计算了图中每一个 Blob。
    • 形状记录: 每次 extract 后，m 中就包含了正确的形状元数据。blobs[top_blob_index].shape = m; 这一句，就是将推演出的形状“偷”回并存入 Blob 的定义中。
  4. 收尾: 再次遍历所有 Layer，将 blobs[...].shape 的信息分别拷贝到 layer->bottom_shapes 和 layer->top_shapes 中，完成“状态重整”。
  5. 限制: 无法处理包含 custom_layer（自定义层）的网络（因为无法推断其形状），也无法处理 Input 层未定义输入形状的网络。

1. 动机：为何在优化后需要“形状推断”？

图优化 Pass（如 fuse_convolution_activation）会修改 Layer 的参数和连接关系，但并不会（也不应该）实时更新所有受影响的中间 Blob 的形状元数据。这会导致 Net 对象内部的 blobs 列表中的 shape 信息变得陈旧。

shape_inference 的目的就是解决这个“陈旧”问题。它在所有优化基本完成后，发起一次“总清算”，通过模拟推理，强制 Extractor 重新计算出所有 Blob 的实际形状，并回写到 blobs 列表中。

2. 准备工作：重置管线与喂入“种子”

在推演开始前，必须进行两项准备工作：

a) 重置层管线 (Recreate Pipeline):

1
2
3
4
5
6
7
8
9
// recreat layer pipeline for param and weight changes
for (size_t i = 0; i < layer_count; i++)
{
    ncnn::Layer* layer = layers[i];

    layer->destroy_pipeline(opt); // 销毁旧状态 (如旧的 weight_data_tm)
    int cret = layer->create_pipeline(opt); // 根据新参数 (如 activation_type) 重新初始化
    // ... 错误检查 ...
}

分析: 这是至关重要的一步。create_pipeline 不仅分配资源，它还负责执行像权重打包 (convolution_transform_kernel_...) 这样的预处理。由于之前的 Pass（如 fuse_convolution_activation）已经修改了 Layer 的参数，必须调用 create_pipeline 来强制 Layer 根据其新的参数（例如，现在它知道自己需要融合激活）来重新生成其内部的优化状态（例如，选择一个支持融合激活的微核）。

b) 喂入“假”输入 (Feed Dummy Inputs):

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
ncnn::Extractor ex = create_extractor();
ex.set_light_mode(true);

// 准备 Input blobs
for (size_t i = 0; i < layer_count; i++)
{
    // ... 找到 Input 层 ...
    if (dims == 0)
    {
        // 输入形状必须已知，否则无法开始推断
        fprintf(stderr, "Input layer %s without shape info, shape_inference skipped\n", layer->name.c_str());
        return -1;
    }

    ncnn::Mat m; // 创建一个“假” Mat
    if (dims == 1) m.create(w);
    if (dims == 2) m.create(w, h);
    if (dims == 3) m.create(w, h, c);

    ex.input(layer->tops[0], m); // 将“假” Mat 注入 Extractor
}
// ... 同样处理已定义形状的 blob (如 MemoryData) ...

分析: 形状推断必须有一个起点。Input 层就是这个起点。shape_inference 要求 Input 层必须有确定的形状。它创建 Mat 对象，但不关心其内容（m 未初始化），只关心其形状元数据。ex.input() 将这个带形状的“种子”放入 Extractor 的 blob_mats 缓存中。

3. 核心：遍历 extract 触发全图“沙盘推演”

这是 shape_inference 最核心的部分，它巧妙地利用了 Extractor 的惰性求值机制。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
fprintf(stderr, "shape_inference\n");

// resolve all layer output blob shape
for (size_t i = 0; i < layer_count; i++)
{
    const ncnn::Layer* layer = layers[i];
    if (layer->type == "ncnnfused")
        continue; // 跳过已无效的层

    // 遍历该层的所有输出
    for (size_t j = 0; j < layer->tops.size(); j++)
    {
        int top_blob_index = layer->tops[j];

        ncnn::Mat m;
        // 核心：触发计算
        ex.extract(top_blob_index, m);

        // 关键：将推演出的形状信息 "偷" 回并保存
        blobs[top_blob_index].shape = m;
    }
}

分析:

  • ex.extract(top_blob_index, m): 如同我们在第 XXXIV 篇中分析的，extract 会检查 blob_mats 缓存。如果 top_blob_index 对应的 Mat 尚未计算 (dims == 0)，它会递归地调用 forward_layer 计算其所有依赖。
  • 为何遍历所有层？: 简单地 extract 最终的 output blob 只会计算图的主干路径。而通过 for 循环强制对网络中的每一个 blob 都执行一次 extract，shape_inference 确保了即使是那些“分支”上的、未被最终输出所依赖的 blob，也会被计算，从而获取其形状。
  • blobs[top_blob_index].shape = m;: extract 返回后，m 虽然内容是无意义的（因为输入是“假”数据），但它的形状元数据（m.dims, m.w, m.h, m.c, m.elempack）是 Extractor 根据各层 forward 逻辑精确推导出来的。此行代码就是将这个“战利品”保存回 Net 对象的 blobs 列表中。

4. 收尾：将形状信息回填给 Layer

在 blobs 列表中的 shape 信息被完全更新后，最后一步是将其同步回 Layer 对象自身。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
// assign all layer blob shape
for (size_t i = 0; i < layer_count; i++)
{
    ncnn::Layer* layer = layers[i];
    if (layer->type == "ncnnfused") continue;

    // 回填 bottom_shapes
    layer->bottom_shapes.resize(layer->bottoms.size());
    for (size_t j = 0; j < layer->bottoms.size(); j++)
    {
        layer->bottom_shapes[j] = blobs[layer->bottoms[j]].shape;
    }

    // 回填 top_shapes
    layer->top_shapes.resize(layer->tops.size());
    for (size_t j = 0; j < layer->tops.size(); j++)
    {
        layer->top_shapes[j] = blobs[layer->tops[j]].shape;
    }
}

分析: Layer 对象内部也维护了 bottom_shapes 和 top_shapes 列表。这一步将 blobs（全局信息源）中的权威形状数据，拷贝到每个 Layer 的本地缓存中。这使得 Layer 自身“知道”了其输入输出的形状，这对于 Net::save 函数正确写入 .param 文件至关重要。

5. 结语

shape_inference 并非一个优化 Pass，而是一个在优化流水线中承前启后的工具 Pass。它在图结构被修改（fuse/eliminate/replace）和层参数被更新后，通过 destroy/create_pipeline 强制层状态再生，并利用 Extractor 的惰性求值机制进行了一次全图“沙盘推演”。

这次推演的目的，就是为了重新收集和更新网络中每一个 Blob 和 Layer 的形状信息，使 Net 对象的内部状态恢复一致。这为后续的内存估算（estimate_memory_footprint）和模型序列化（save）提供了准确的数据基础，是 ncnnoptimize 流程中不可或缺的“状态重整”环节。