读 ncnn 源码（XIII）：Extractor::extract——触发推理的引擎核心

在之前的篇章中，我们已经探讨了 ncnn 如何加载模型、优化权重布局。当模型准备就绪，输入数据也通过 Extractor::input() 绑定后，最后一步就是调用 Extractor::extract() 来获取我们关心的输出结果。这个看似简单的函数调用，背后却隐藏着 ncnn 推理引擎的核心调度逻辑，涵盖了惰性求值、后端分发、内存管理等多个关键环节。

本篇，我们将深入 Extractor::extract 及其调用的 NetPrivate::forward_layer、NetPrivate::do_forward_layer 函数，剖析 ncnn 是如何根据依赖关系，按需、高效地执行神经网络推理的。

TL;DR

  1. Extractor::extract 是推理触发器: 它并非立即计算整个网络，而是采用惰性求值 (Lazy Evaluation)。只有当用户请求某个特定的 blob 时，才会触发计算该 blob 所需的最少网络路径。
  2. 核心流程: extract(blob_name/index) 首先检查 blob_mats 缓存中是否已存在目标 blob。若不存在，则找到其生产者层 (Producer Layer)，并调用 NetPrivate::forward_layer 来执行计算。
  3. forward_layer 的递归机制: 这是实现惰性求值的关键。它采用深度优先搜索 (DFS) 的方式：要计算第 L 层，必须先递归调用 forward_layer 确保其所有输入 blob（即 layer->bottoms）都已被计算出来。
  4. do_forward_layer 实际执行: 当某一层的所有输入 blob 都准备就绪后，do_forward_layer 负责调用该层真正的 forward 或 forward_inplace 方法。它还处理内存优化 (Light Mode)、数据布局转换 (Layout Conversion) 以及结果缓存。
  5. 后端分发与后处理: extract 内部包含对 Vulkan 后端的分发逻辑。在获取到计算结果后，还会进行必要的布局转换（Packed -> Plain）、数据类型转换（FP16/BF16/INT8 -> FP32）以及内存池解绑 (clone)，确保返回给用户的 Mat 是标准的、独立的 FP32 数据。

1. Extractor::extract：用户接口与惰性求值触发

Extractor 对象是用户与 ncnn 网络交互的主要接口。input() 负责设置输入，而 extract() 则负责获取输出。

1.1 函数重载

Extractor 提供了两个 extract 重载版本：

  • extract(const char* blob_name, Mat& feat, int type = 0): 通过 blob 的名称来提取。内部会先调用 net->find_blob_index_by_name() 找到对应的索引。
  • extract(int blob_index, Mat& feat, int type = 0): 通过 blob 的索引来提取。这是核心实现。

type 参数用于指定期望的输出数据类型和布局（暂不深入）。

1.2 惰性求值核心逻辑 (extract by index)

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
int Extractor::extract(int blob_index, Mat& feat, int type)
{
    // ... 省略边界检查和环境设置 ...

    // 核心：检查 blob_mats 缓存
    if (d->blob_mats[blob_index].dims == 0) // dims == 0 表示尚未计算
    {
        // 找到该 blob 的生产者层
        int layer_index = d->net->blobs()[blob_index].producer;

        // 设置本地内存池 (如果启用且未提供全局的)
        if (d->opt.use_local_pool_allocator) { /* ... setup local allocators ... */ }

#if NCNN_VULKAN
        if (d->opt.use_vulkan_compute) // 检查是否启用 Vulkan 后端
        {
            // Vulkan 后端执行路径
            ncnn::VkCompute cmd(...);
            VkMat feat_gpu;
            // 递归调用 Vulkan 的 extract (内部逻辑类似，但操作 GPU 对象)
            ret = extract(blob_index, feat_gpu, cmd);
            if (ret == 0 && ...) {
                // 如果 Vulkan 计算完成，下载结果到 CPU Mat
                cmd.record_download(feat_gpu, d->blob_mats[blob_index], d->opt);
                ret = cmd.submit_and_wait();
                // ... benchmark code ...
            }
        }
        else // Vulkan 未启用或执行失败，回退到 CPU
        {
            // 触发 CPU 计算路径
            ret = d->net->d->forward_layer(layer_index, d->blob_mats, d->opt);
        }
#else // 没有 Vulkan 支持，直接走 CPU 路径
        // 触发 CPU 计算路径
        ret = d->net->d->forward_layer(layer_index, d->blob_mats, d->opt);
#endif // NCNN_VULKAN
    } // end if (blob needs computation)

    // 从缓存中获取最终结果 (可能刚刚被计算出来)
    feat = d->blob_mats[blob_index];

    // ... 后处理：布局转换、类型转换、内存池解绑 ...

    // ... 恢复环境设置 ...
    return ret;
}

关键点解析:

  • d->blob_mats: 这是 Extractor 内部维护的一个 std::vector<Mat>，其大小与网络中的 blob 数量相同。它充当了计算结果的缓存。初始状态下，所有 Mat 的 dims 都为 0。
  • if (d->blob_mats[blob_index].dims == 0): 这是惰性求值的判断核心。只有当请求的 blob 尚未计算时，才会进入计算流程。如果之前某个 extract 调用已经计算过这个 blob（或其依赖路径上的中间 blob），它的 dims 将不为 0，可以直接从缓存返回。
  • 后端分发: 代码清晰地展示了优先尝试 Vulkan 后端（如果启用），失败或未启用则回退到 CPU 执行 (forward_layer) 的逻辑。
  • forward_layer 调用: 这是将计算任务委托给 Net 内部实现的关键一步，标志着递归计算的开始。

1.3 结果后处理

在获取到 blob_mats[blob_index] 后，extract 还会进行一系列重要的后处理，确保返回给用户的 feat 是标准的 FP32、非 Packed、且内存独立的数据：

  • 布局转换: 如果网络内部使用了 Packed Layout（feat.elempack != 1）以利用 SIMD，而用户未指定保留 (type == 0)，则调用 convert_packing 将其转回 Plain Layout。
  • 数据类型转换: 如果网络内部使用了低精度存储（FP16, BF16, INT8），而用户未指定保留 (type == 0)，则调用 cast_..._to_float32 将其转回 FP32。
  • 内存池解绑: 如果使用了本地内存池 (use_local_pool_allocator)，且 feat 仍然在使用该池的内存，则调用 feat = feat.clone() 进行深拷贝。这使得用户可以在获取 feat 后安全地销毁 Net 对象，而不会导致 feat 的内存失效。

2. NetPrivate::forward_layer：递归的依赖解析器

这个函数是 CPU 推理路径的核心调度器，它完美体现了递归和依赖解析的思想。

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
int NetPrivate::forward_layer(int layer_index, std::vector<Mat>& blob_mats, const Option& opt) const
{
    const Layer* layer = layers[layer_index]; // 获取要计算的目标层

    // 核心：递归确保所有输入 blob 都已就绪
    for (size_t i = 0; i < layer->bottoms.size(); i++)
    {
        int bottom_blob_index = layer->bottoms[i]; // 获取一个输入 blob 的索引

        if (blob_mats[bottom_blob_index].dims == 0) // 检查该输入 blob 是否已计算
        {
            // 如果未计算，递归调用 forward_layer 计算其生产者层
            int ret = forward_layer(blobs[bottom_blob_index].producer, blob_mats, opt);
            if (ret != 0) // 如果上游计算出错，则中断返回
                return ret;
        }
    }

    // 所有输入 blob 都已就绪，现在执行当前层的计算
#if NCNN_BENCHMARK
    double start = get_current_time();
    // ... benchmark 相关代码 ...
#endif
    int ret = 0;
    if (layer->featmask) // 特征掩码支持 (暂不深入)
    {
        ret = do_forward_layer(layer, blob_mats, get_masked_option(opt, layer->featmask));
    }
    else
    {
        // 调用实际执行函数
        ret = do_forward_layer(layer, blob_mats, opt);
    }
#if NCNN_BENCHMARK
    double end = get_current_time();
    // ... benchmark 相关代码 ...
#endif
    if (ret != 0)
        return ret;

    return 0; // 当前层计算成功
}

关键点解析:

  • 递归调用: for 循环检查当前层 layer 的所有输入 blob (layer->bottoms)。如果发现某个输入 blob 尚未计算 (dims == 0)，它不会立即计算当前层，而是递归地调用 forward_layer 去计算那个输入 blob 的生产者层。
  • DFS 顺序: 这个递归过程本质上是在网络的计算图（DAG）上进行深度优先搜索。只有当一个节点（层）的所有前驱节点（产生其输入的层）都已被访问（计算）后，该节点才会被访问（计算）。
  • do_forward_layer 调用: 当循环结束，意味着当前层的所有输入都已在 blob_mats 中准备就绪。此时，才调用 do_forward_layer 来执行当前层的实际计算。

3. NetPrivate::do_forward_layer：层计算执行与内存管理

这个函数负责调用具体 Layer 对象的 forward 或 forward_inplace 方法，并处理相关的内存优化。

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
52
53
54
55
56
57
58
59
60
int NetPrivate::do_forward_layer(const Layer* layer, std::vector<Mat>& blob_mats, const Option& opt) const
{
    if (layer->one_blob_only) // 处理单输入单输出的层
    {
        int bottom_blob_index = layer->bottoms[0];
        int top_blob_index = layer->tops[0];
        Mat& bottom_blob_ref = blob_mats[bottom_blob_index];
        Mat bottom_blob; // 可能的拷贝

        // Light Mode 内存优化：如果支持 inplace 且输入被共享，则 clone 输入
        if (opt.lightmode && layer->support_inplace && *bottom_blob_ref.refcount != 1)
        {
            bottom_blob = bottom_blob_ref.clone(opt.blob_allocator);
            // ... 错误检查 ...
        }
        if (bottom_blob.dims == 0) bottom_blob = bottom_blob_ref; // 否则直接引用

        // 可选的布局转换 (例如 packed -> plain)
        int ret = convert_layout(bottom_blob, layer, opt);
        if (ret != 0) return ret;

        // 执行计算
        if (opt.lightmode && layer->support_inplace) // 优先尝试 inplace
        {
            Mat& bottom_top_blob = bottom_blob;
            ret = layer->forward_inplace(bottom_top_blob, opt);
            // ... 错误检查 ...
            blob_mats[top_blob_index] = bottom_top_blob; // 结果直接写入缓存
        }
        else // 不支持 inplace 或 lightmode 未开启
        {
            Mat top_blob;
            ret = layer->forward(bottom_blob, top_blob, opt);
            // ... 错误检查 ...
            blob_mats[top_blob_index] = top_blob; // 结果写入缓存
        }

        // Light Mode 内存优化：计算完成后释放输入 blob
        if (opt.lightmode) blob_mats[bottom_blob_index].release();
    }
    else // 处理多输入多输出的层 (逻辑类似)
    {
        // ... 获取多个 bottom_blobs (同样处理 light mode clone) ...
        // ... 逐个 convert_layout ...
        if (opt.lightmode && layer->support_inplace)
        {
            ret = layer->forward_inplace(bottom_blobs, opt);
            // ... 存储多个 top_blobs 到 blob_mats ...
        }
        else
        {
            std::vector<Mat> top_blobs(layer->tops.size());
            ret = layer->forward(bottom_blobs, top_blobs, opt);
            // ... 存储多个 top_blobs 到 blob_mats ...
        }
        // Light Mode 内存优化：释放所有输入 blobs
        if (opt.lightmode) { /* ... release all bottom blobs ... */ }
    }
    return 0;
}

关键点解析:

  • one_blob_only 分支: 针对单输入输出和多输入输出的层，分别处理输入 Mat 的获取。
  • Light Mode 优化 (opt.lightmode):
    • In-place 计算: 如果层支持 forward_inplace 且开启了 Light Mode，优先调用原地计算版本，直接在输入 Mat 的内存上写入结果，避免了为输出 Mat 分配新内存。
    • Clone on Shared: 如果输入 Mat 的引用计数 *refcount 大于 1（表示有其他地方也在用这块内存），即使支持 in-place，也必须先 clone 一份副本，以防原地修改影响其他使用者。
    • 输入释放: 计算完成后，调用 blob_mats[bottom_blob_index].release() 释放输入 Mat 的引用。如果引用计数变为 0，内存会被回收。这是 Light Mode 显著降低峰值内存占用的关键。
  • convert_layout: 在调用 forward 之前，可能需要根据层的要求（例如，某些层只能处理 Plain Layout）对输入 Mat 进行布局转换。
  • forward / forward_inplace 调用: 这是真正执行该层计算的地方。
  • 结果缓存: 将计算得到的 top_blob(s) 存回 blob_mats 向量中，供后续层使用或最终由 extract 返回。

4. 结语

Extractor::extract 并非一个简单的函数调用，而是 ncnn 推理引擎的大门。它巧妙地利用 blob_mats 作为缓存，结合 forward_layer 的递归依赖解析，实现了高效的惰性求值。do_forward_layer 则负责具体的层执行和精细的内存管理策略（尤其是 Light Mode）。整个流程清晰地展现了 ncnn 在追求高性能的同时，对内存占用、计算效率和后端异构性的周全考虑，是理解其核心工作原理的关键环节。

该封面图片由Rezső Terbe在Pixabay上发布