ncnn 内存池机制详解：如何通过自定义 Allocator 解决频繁分配导致的推理抖动

背景

在嵌入式设备或 Android/iOS 开发中，AI 模型的推理性能不仅取决于算法复杂度，还深受系统资源调度的影响。很多开发者会发现，ncnn 在连续推理时，由于系统底层的 malloc 和 free 导致内存碎片或内核锁竞争，从而引发推理时间的“随机抖动”。

ncnn 默认提供的 PoolAllocator 已经极大缓解了这一问题，但在一些对实时性要求极高的场景（如自动驾驶、高频工业视觉），我们需要更彻底的控制手段：自定义内存分配器。

ncnn 内存分配机制概览

ncnn 的内存管理是通过 ncnn::Allocator 抽象类实现的。主要涉及两个核心成员：
1. blob_allocator：用于存放网络中间层的特征图（Tensor）数据。
2. workspace_allocator：用于算子内部的临时计算空间。

默认情况下，ncnn 使用 PoolAllocator。如果我们要消除波动，可以实现一个简单的静态分配器（Static Allocator），预先占领一块物理内存。

核心代码：实现一个固定缓冲区分配器

下面的代码展示了如何通过封装一个预分配的 buffer 来实现自定义分配器，从而彻底规避运行时向系统申请内存。

#include "net.h"
#include "allocator.h"
#include <iostream>

// 自定义静态内存分配器
class FixedBufferAllocator : public ncnn::Allocator {
public:
    FixedBufferAllocator(size_t capacity) : capacity(capacity), offset(0) {
        // 预先申请一大块对齐内存
        data = (unsigned char*)ncnn::fastMalloc(capacity);
    }

    ~FixedBufferAllocator() {
        ncnn::fastFree(data);
    }

    // 重写分配逻辑
    virtual void* fastMalloc(size_t n) {
        // 对齐到 16 字节
        size_t aligned_n = (n + 15) & ~15;
        if (offset + aligned_n > capacity) {
            std::cerr << "Error: Out of memory in FixedBufferAllocator" << std::endl;
            return 0;
        }
        void* ptr = data + offset;
        offset += aligned_n;
        return ptr;
    }

    // 在此场景下，单次推理不单独释放，由 reset 统一清理
    virtual void fastFree(void* ptr) {}

    void reset() { offset = 0; }

private:
    unsigned char* data;
    size_t capacity;
    size_t offset;
};

实操：如何在推理任务中注入自定义分配器

有了分配器后，我们需要将其注入到 ncnn::Net 的配置中。

void run_inference() {
    ncnn::Net net;

    // 1. 创建自定义分配器（假设模型推理最高需要 64MB）
    FixedBufferAllocator my_blob_alloc(64 * 1024 * 1024);
    FixedBufferAllocator my_workspace_alloc(16 * 1024 * 1024);

    // 2. 配置 Option
    ncnn::Option opt;
    opt.blob_allocator = &my_blob_alloc;
    opt.workspace_allocator = &my_workspace_alloc;

    net.opt = opt;

    // 3. 加载模型及推理
    net.load_param("model.param");
    net.load_model("model.bin");

    for (int i = 0; i < 100; ++i) {
        // 推理前重置偏移量，实现内存复用
        my_blob_alloc.reset();
        my_workspace_alloc.reset();

        ncnn::Extractor ex = net.create_extractor();
        ncnn::Mat in(224, 224, 3);
        ex.input("data", in);

        ncnn::Mat out;
        ex.extract("output", out);
        // 处理结果...
    }
}

关键点总结

1. 防止内存碎片：通过一次性 fastMalloc 申请大块空间，后续推理均在内部指针偏移，避免了频繁的系统级 brk 或 mmap 调用。
2. 降低延迟抖动：固定路径的内存分发使得每一帧的开销几乎完全一致。
3. 线程安全注意：上述简单的 FixedBufferAllocator 不是线程安全的。如果在多线程环境并行执行 Extractor，请为每个线程创建独立的 Allocator，或者在分配逻辑中加入 ncnn::MutexLock。

通过这种方式，您可以让 ncnn 在资源受限的国产芯片或低端 Android 设备上跑出更加平稳的性能曲线。