📚 课程目录第4章内存管理与复用

4. 内存管理与复用

内存管理，说白了就是端侧推理的命门。你想想看，一个模型在服务器上跑，几百GB显存随便用，但到了手机、IoT设备上，可能总共就几百MB内存，还得跟其他App抢。我这些年调过不少端侧推理引擎，踩过的坑十有八九都跟内存有关。今天咱们就把这块硬骨头啃下来。

为什么内存管理在端侧这么重要？

先看一组数据。一个MobileNetV3模型，参数量大概5.4M，FP32推理需要约21MB。但中间激活值呢？如果输入是224x224，特征图一路算下来，峰值内存可能冲到40-50MB。这还只是一个小模型。要是换成YOLOv5s，峰值内存轻松破百MB。

更麻烦的是，端侧设备的内存带宽通常只有服务器的十分之一甚至更低。频繁的malloc/free不仅慢，还会产生碎片。我见过一个项目，模型在开发板上跑着跑着就OOM了，查了半天，发现是内存碎片导致连续大块分配失败。

核心矛盾：端侧设备既要跑大模型，又要保证实时性，还得省内存。解决思路就三个字——复用。

Tensor内存池设计

内存池，说白了就是提前申请一大块内存，然后自己管分配和释放。这样能避免频繁的系统调用，也方便做复用。

我习惯把内存池分成两层：

全局池：存放模型权重、常量等生命周期长的数据。模型加载时一次性分配，推理过程中不动。
工作池：存放中间激活值、临时buffer。这部分是优化的重点。

工作池的设计，核心是解决两个问题：

什么时候分配？——静态分析还是动态分配？
怎么复用？——哪些Tensor可以共享同一块内存？

静态内存规划

我个人强烈推荐静态规划。在模型编译阶段，通过计算图分析，提前算出每个Tensor的生命周期，然后做内存复用。

具体做法是这样的：

遍历计算图，给每个Tensor打上生存区间标签（从哪个算子产生，到哪个算子不再使用）。
把所有Tensor的生存区间画在时间轴上，这就是一个经典的区间着色问题。
用贪心算法（比如First-Fit）给每个区间分配内存偏移量。

嗯，这里要注意：生存区间的计算必须考虑算子内部的临时buffer。比如卷积操作，im2col展开后需要一块临时内存，这块内存的生存区间就是卷积算子内部。

经验之谈：我早期做静态规划时，只考虑了Tensor的生存区间，结果算子内部的临时buffer跟Tensor冲突，导致内存踩踏。后来我改成把算子内部的临时buffer也当成一个"虚拟Tensor"参与规划，问题就解决了。

动态内存池

有些场景没法做静态规划，比如动态shape的模型、或者推理框架需要支持运行时加载新算子。这时候就得用动态内存池。

动态池的设计要点：

分块管理：把大块内存切成固定大小的块（比如64B、256B、1KB、4KB...），用伙伴算法或slab分配器管理。
延迟回收：Tensor释放后，内存块不立即归还系统，而是留在池里给下一个Tensor用。
水位线：设置一个高水位线，当池中已分配内存超过阈值时，触发整理或扩容。

我曾经在一个项目里遇到过动态池碎片化严重的问题。排查后发现，是因为频繁分配释放不同大小的Tensor，导致小碎片越来越多。后来我改用固定大小块+大块专用区的混合策略，碎片率从15%降到了2%以下。

in-place操作

in-place操作，就是让输出Tensor复用输入Tensor的内存。说白了，就是"用完就扔"，省掉一份内存。

常见的in-place操作：

算子	是否支持in-place	注意事项
ReLU / ReLU6	是	直接修改输入Tensor
BatchNorm	是	需要保证后续算子不再使用原始输入
Add (element-wise)	部分支持	只能复用其中一个输入，另一个必须保留
Concat	否	输出需要连续内存，无法复用输入
Reshape / Transpose	视情况	如果只是改变视图，可以in-place；如果涉及数据重排，不行

为什么有些算子不能in-place？举个例子，Concat操作需要把多个输入拼接到一起，输出是一块连续内存。如果复用其中一个输入，那其他输入的数据就没地方放了。所以Concat必须分配新内存。

避坑指南：我曾经在实现in-place ReLU时，没考虑到ReLU后面跟着一个需要原始输入的算子（比如残差连接）。结果ReLU把输入改了，残差加的时候拿到的全是ReLU后的值，模型精度直接崩了。所以做in-place优化前，一定要确认所有引用该Tensor的算子都不再需要原始值。

内存对齐策略

内存对齐，听起来是个小细节，但影响很大。不对齐的内存访问，在ARM上可能慢3-5倍，在x86上甚至会导致段错误。

端侧推理中，常见的对齐要求：

SIMD指令集：NEON要求16字节对齐，SSE要求16字节对齐，AVX要求32字节对齐。
DMA传输：很多NPU的DMA要求32字节或64字节对齐。
缓存行：ARM Cortex-A系列缓存行通常是64字节，对齐访问能避免缓存行分裂。

我的做法是：

统一对齐粒度：所有Tensor的起始地址都按最大对齐要求对齐。比如目标设备NEON和DMA都要求16字节对齐，那就统一按16字节对齐。
填充padding：每个Tensor的末尾也填充到对齐边界，这样下一个Tensor的起始地址自然对齐。
池内对齐：内存池分配时，返回的地址保证对齐。如果池内剩余空间不够对齐，就跳过这段空间（产生内部碎片）。

举个例子，假设对齐粒度是16字节，要分配一个大小为30字节的Tensor：

// 分配函数示例
void* pool_alloc(size_t size, size_t alignment) {
    // 计算对齐后的实际分配大小
    size_t aligned_size = (size + alignment - 1) & ~(alignment - 1);
    
    // 检查当前偏移是否对齐
    size_t current_offset = pool->offset;
    size_t aligned_offset = (current_offset + alignment - 1) & ~(alignment - 1);
    
    // 如果不对齐，跳过对齐空间
    if (aligned_offset != current_offset) {
        pool->offset = aligned_offset;
    }
    
    // 分配
    void* ptr = pool->base + pool->offset;
    pool->offset += aligned_size;
    return ptr;
}

这段代码里，(size + alignment - 1) & ~(alignment - 1) 是经典的对齐取整公式。嗯，这里要注意：如果池内剩余空间不够对齐，就得考虑换一块池或者扩容。

一个小技巧：我习惯在内存池初始化时，把池的起始地址也按最大对齐要求对齐。这样后续分配时，大部分情况下偏移量天然就是对齐的，能减少内部碎片。

整体架构

下面这张图展示了内存管理的整体流程：

总结一下

内存管理这块，说白了就是三个关键词：复用、对齐、规划。

复用靠内存池和in-place操作，减少分配次数和峰值内存。
对齐保证硬件访问效率，避免性能陷阱。
规划（静态或动态）决定内存怎么分、怎么还。

我见过不少团队，模型精度调得不错，但一上端侧就卡死或OOM。十有八九是内存管理没做好。其实只要把上面这几步做到位，大部分问题都能解决。嗯，今天就先聊到这儿，下次咱们接着讲算子优化。