NCNN 自定义内存分配器注入与共享内存池实战消除多模型切换时的碎片化与 OOM一、模型切换时 malloc/free 循环导致的碎片化死亡螺旋NCNN 是腾讯开源的高性能神经网络推理框架专为移动端和嵌入式平台优化。其核心设计理念之一是零依赖——不依赖任何第三方计算库全部使用手写汇编优化。在内存管理方面NCNN 的默认行为是在模型加载时通过ncnn::Allocator接口分配权重内存和工作内存working memory在模型销毁时释放。这个创建即分配、销毁即释放的模式在单模型场景下工作良好。但在多模型切换的场景中——比如一个应用需要在人脸检测模型RetinaFace和人脸识别模型ArcFace之间切来切去——每次切换都涉及旧模型的内存释放和新模型的内存分配。在 Linux 的 glibc malloc 实现中嵌入式平台常见使用 dlmalloc 或 jemalloc反复的分配60MB→ 释放 → 分配40MB→ 释放 → 分配80MB操作会在堆上产生大量不可合并的碎片。典型症状系统当前剩余 150MB 连续物理内存但三个模型分别需要 60MB、40MB、80MB由于内存碎片化即使总空闲 150MB也无法满足任何一个模型的大块连续分配。这就是内存碎片导致的虚假 OOM——不是内存不够而是没有足够大的连续块。解决思路是使用共享内存池预先从系统分配一块大内存如 256MB所有模型从这个池中通过 NCNN 自定义 Allocator 分配权重和工作内存。模型切换时旧模型归还内存到池中新模型从池的另一个区域分配——由于所有分配都在同一块连续地址空间内不会有 malloc 层面的外部碎片。二、NCNN 内存分配架构与自定义注入点NCNN 的内存分配分为两个阶段模型加载阶段ncnn::Net::load_param()和ncnn::Net::load_model()读取模型文件和权重数据。权重数据blob在此时通过Allocator分配内存。推理执行阶段每一层的forward()方法通过ncnn::WorkspaceAllocator工作空间分配器申请临时内存用于存储中间张量。工作空间在推理结束后自动回收。sequenceDiagram participant App as 应用层 participant Net as ncnn::Net participant Alloc as VkAllocator / PoolAllocator participant Workspace as ncnn::WorkspaceAllocator participant OS as OS malloc / mmap Note over App,OS: 自定义注入流程 App-OS: 初始化: 分配 256MB 连续内存池 OS--App: uint8_t* pool_ptr App-Alloc: 创建 PoolAllocator(pool_ptr, 256MB) App-Net: net.opt.use_custom_allocator true App-Net: net.set_custom_allocator(pool_allocator) Note over App,OS: 模型A加载 App-Net: load_model(model_a.param) Net-Alloc: fastMalloc(60MB) → 从池中分配 Alloc--Net: offset0, size60MB Net-Alloc: fastMalloc(20MB) → 工作空间 Alloc--Net: offset60MB, size20MB Note over App,OS: 模型A推理 App-Net: input(data, in) Net-Workspace: get_workspace(10MB) → 从池分配 Workspace--Net: 临时工作区 Note over App,OS: 切换到模型B: 先清空池, 再加载 App-Net: clear() 释放模型A App-Alloc: reset() → offset0 重置池状态 App-Net: load_model(model_b.param) Net-Alloc: fastMalloc(40MB) → offset0 Alloc--Net: 分配成功(池足够大)NCNN 提供了两个关键的自定义注入接口ncnn::Allocator虚基类需要实现两个纯虚函数fastMalloc(size_t size)和fastFree(void* ptr)。所有权重的内存分配都通过这个接口。ncnn::VkAllocatorVulkan 专用用于 GPU 内存的分配接口类似但操作的是 Vulkan 设备内存。三、共享内存池分配器的完整实现以下是基于内存池的自定义 Allocator 实现支持单模型快速重置和多模型并行运行。/** * pool_allocator.h - NCNN 自定义内存池分配器 * * 设计目标: * 1. 预分配一块大内存所有模型从池中分配消除外部碎片 * 2. 采用 bump pointer(碰撞指针) 策略: 用一个原子偏移量表示分配边界 * fastMalloc 将偏移前移fastFree 不回收(等 reset 统一清空) * 3. 支持 reset() 清空整个池用于模型切换时的快速重用 * 4. 线程安全: 使用原子操作保证并发的 fastMalloc 安全 * * 权衡: bump allocator 简单高效, 但无法释放中间块。 * 对于推理场景, 这正好符合语义——模型加载是一次性分配, * 推理工作空间在每层完成后回收(由 NCNN 内部管理), * 模型切换时全局 reset 池即可。 */ #pragma once #include ncnn/allocator.h #include atomic #include cstdint #include cstdlib #include cstring class PoolAllocator : public ncnn::Allocator { public: /** * 构造函数 * param pool_size 内存池总大小(字节), 应 所有模型峰值内存之和 * 典型值: 单模型 60-100MB, 双模型并行 160-200MB * param alignment 内存对齐要求(字节), NCNN 内部要求至少 16 字节对齐 * 推荐 64 或 128 以适配 NEON/AVX 向量化访问 */ explicit PoolAllocator(size_t pool_size, size_t alignment 128) : alignment_(alignment) { // 使用 posix_memalign 分配对齐内存 // 注意: Windows 平台使用 _aligned_malloc int ret posix_memalign((void**)pool_base_, alignment_, pool_size); if (ret ! 0 || !pool_base_) { fprintf(stderr, [PoolAllocator] Failed to allocate %zu bytes, errno%d\n, pool_size, ret); pool_base_ nullptr; pool_size_ 0; return; } pool_size_ pool_size; memset(pool_base_, 0, pool_size_); // 可选: 清零(方便调试) offset_.store(0, std::memory_order_release); } ~PoolAllocator() override { if (pool_base_) { free(pool_base_); pool_base_ nullptr; pool_size_ 0; offset_.store(0); } } // 禁止拷贝(内部管理内存所有权) PoolAllocator(const PoolAllocator) delete; PoolAllocator operator(const PoolAllocator) delete; /** * 自定义内存分配 * 使用 fetch_add 原子操作实现无锁的 bump allocation * * param size 请求分配大小(字节), NCNN 在加载时传入 * return 分配的内存指针, 失败返回 nullptr * * 注意: NCNN 内部不保证 fastMalloc 返回的内存被 fastFree 配对调用。 * 因此 bump allocator 的只增不减策略是安全的—— * NCNN 在模型 clear() 时会调用所有 fastFree, 但 bump allocator * 不依赖配对, 只在 reset() 时统一回卷 offset。 */ void* fastMalloc(size_t size) override { if (!pool_base_ || size 0) return nullptr; // 向上取整到 alignment 边界 size_t aligned_size (size alignment_ - 1) ~(alignment_ - 1); // 原子获取并增加偏移量 size_t current_offset offset_.fetch_add(aligned_size, std::memory_order_relaxed); // 溢出检查: 分配后偏移不能超过池大小 if (current_offset aligned_size pool_size_) { // 回退偏移量(部分回退忽略竞争) // 在生产代码中应使用更严格的回退策略 offset_.fetch_sub(aligned_size, std::memory_order_relaxed); fprintf(stderr, [PoolAllocator] OOM: requested%zu, offset%zu, pool%zu\n, size, current_offset, pool_size_); return nullptr; } return pool_base_ current_offset; } /** * 释放内存(Bump Allocator 语义下为空操作) * 不使用标准的 free 语义内存生命周期由 reset() 管理 */ void fastFree(void* /*ptr*/) override { // bump allocator: 不做实际释放 // 内存统一在 reset() 时回收 } /** * 清空内存池 * 将分配偏移重置为零复用同一块物理内存 * 用于模型切换时快速清理旧模型的内存占用 * * 线程安全注意: reset() 必须在确认没有线程正在使用池内存时调用。 * 典型调用时机: 旧模型 clear() 完成后、新模型 load_model() 之前。 */ void reset() { offset_.store(0, std::memory_order_release); // 可选: 清零前若干个字节以辅助调试 // memset(pool_base_, 0, std::min((size_t)4096, pool_size_)); } /** * 获取当前内存使用量(字节) */ size_t used_bytes() const { return offset_.load(std::memory_order_relaxed); } /** * 获取内存池总大小(字节) */ size_t total_bytes() const { return pool_size_; } /** * 计算当前利用率(百分比) */ float utilization() const { if (pool_size_ 0) return 0.0f; return (float)used_bytes() / (float)pool_size_ * 100.0f; } private: uint8_t* pool_base_ nullptr; size_t pool_size_ 0; size_t alignment_ 128; std::atomicsize_t offset_{0}; };使用示例多模型切换流程/** * multi_model_inference.cpp - 多模型切换使用共享内存池 * * 场景: 人脸检测(RetinaFace) → 人脸关键点(WingLoss) → 人脸识别(ArcFace) * 三个模型共享 200MB 内存池依次串行加载和推理 */ #include ncnn/net.h #include cstdio int main() { // 1. 创建共享内存池: 200MB // 实测: RetinaFace ~55MB, WingLoss ~30MB, ArcFace ~70MB // 峰值 ~70MB ( 200MB), 预留余量 130MB PoolAllocator shared_pool(200 * 1024 * 1024, 64); if (shared_pool.total_bytes() 0) { fprintf(stderr, Failed to create memory pool!\n); return -1; } printf(Memory pool: %.1f MB allocated\n, shared_pool.total_bytes() / (1024.0 * 1024.0)); // 2. 配置 NCNN 使用自定义分配器 ncnn::Net net; net.opt.use_vulkan_compute false; // CPU 推理 // 注入自定义分配器: 关键属性设置 net.opt.blob_allocator shared_pool; // 权重 blob 内存 net.opt.workspace_allocator shared_pool; // 工作空间内存 // 3. 阶段一: 人脸检测 printf(--- Loading RetinaFace ---\n); if (net.load_param(retinaface.param) ! 0) { fprintf(stderr, Failed to load retinaface.param\n); return -1; } if (net.load_model(retinaface.bin) ! 0) { fprintf(stderr, Failed to load retinaface.bin\n); return -1; } printf(Pool usage after RetinaFace: %.1f MB\n, shared_pool.used_bytes() / (1024.0 * 1024.0)); // 推理省略数据准备 // ncnn::Mat in ...; // ncnn::Extractor ex net.create_extractor(); // ex.input(data, in); // ncnn::Mat out; // ex.extract(output, out); // 4. 切换到关键点检测模型 printf(--- Switching to WingLoss ---\n); net.clear(); // 清理模型 A 的状态 shared_pool.reset(); // 重置内存池(关键步骤!) printf(Pool usage after reset: %.1f MB\n, shared_pool.used_bytes() / (1024.0 * 1024.0)); if (net.load_param(wingloss.param) ! 0) { fprintf(stderr, Failed to load wingloss.param\n); return -1; } if (net.load_model(wingloss.bin) ! 0) { fprintf(stderr, Failed to load wingloss.bin\n); return -1; } printf(Pool usage after WingLoss: %.1f MB\n, shared_pool.used_bytes() / (1024.0 * 1024.0)); // 5. 切换到人脸识别模型 printf(--- Switching to ArcFace ---\n); net.clear(); shared_pool.reset(); if (net.load_param(arcface.param) ! 0) { fprintf(stderr, Failed to load arcface.param\n); return -1; } if (net.load_model(arcface.bin) ! 0) { fprintf(stderr, Failed to load arcface.bin\n); return -1; } printf(Pool usage after ArcFace: %.1f MB\n, shared_pool.used_bytes() / (1024.0 * 1024.0)); printf(Pool utilization: %.1f%%\n, shared_pool.utilization()); printf(All models loaded successfully with shared pool.\n); return 0; }四、内存池方案的边界条件与工程取舍Bump Allocator 的只增不减语义Bump Allocator碰撞指针分配器的核心特性是内存只分配不回收。这简化了实现无需维护 free list但带来了池利用率可能下降的问题。在三种情况下这个问题会显现模型加载时分配了很多小对象bump allocator 为每个小对象都做对齐填充每分配一次浪费 alignment - (size % alignment) 字节。解决方案是使用较小的 alignment如 64或改为 slab allocator。模型卸载后立即加载另一个模型如果没有reset()旧模型的内存不会被回收池的可用空间逐渐缩小。必须确保在net.clear()后调用shared_pool.reset()。多模型并行运行如果两个模型同时驻留在内存中例如视频流的检测识别交错执行bump allocator 无法支持释放模型 A 的内存给模型 C 用——因为模型 B 分配的内存位于池中间阻断了回收。这种场景需要使用更复杂的 free-list 或 buddy 分配器。与 Vulkan GPU 内存的交互当使用 NCNN 的 Vulkan 后端时GPU 内存由ncnn::VkAllocator管理与 CPU 内存池是独立的。Vulkan 设备内存的分配比 CPU 内存更昂贵涉及 GPU 页表映射且有严格的对齐要求通常 256 字节。如果同时使用 CPU 和 GPU 推理建议为 GPU 也实现一个类似的池分配器基于ncnn::VkAllocator接口。内存池大小选取的实践经验池大小设定有两种策略峰值覆盖策略pool_size max(model_A_size, model_B_size, ...)适用串行模型切换。优点是池最小化缺点是添加新模型时可能不足。总和覆盖策略pool_size sum(all_model_sizes)适用多模型并行。优点是安全缺点是在串行场景下浪费了大量内存。在嵌入式平台1GB DRAM上的经验值单模型运行时池设为 120MB 足够覆盖常见视觉模型MobileNet/YOLO-Nano/RetinaFace 等双模型并行时池设为 200MB 可覆盖轻量级组合。五、总结NCNN 自定义内存分配器的注入解决的是嵌入式多模型推理中反复 malloc/free 导致的内存碎片问题本质是通过空间换确定性策略——用一块预分配的内存池替代 OS 堆分配器的动态管理。工程实践要点Bump Allocator 适配 NCNN 语义NCNN 的fastMalloc/fastFree支持只增不减的 bump 策略通过reset()实现池的快速重用。注入方式设置net.opt.blob_allocator和net.opt.workspace_allocator为自定义分配器指针。池大小规划串行模型切换使用峰值覆盖策略并行推理使用总和覆盖策略。线程安全使用std::atomic的fetch_add实现无锁 bump allocation无需互斥锁。当前方案的局限在于单池不支持多模型并行时的部分释放——如果需要更复杂的内存管理如按模型粒度回收需要实现基于 free-list 的分配器。但对于绝大多数边缘端串行推理场景bump allocator shared pool 的组合已经能在 100 行代码内解决碎片问题。