TFLite内存规划器详解：复用张量按生命周期压峰值内存

做端侧推理的兄弟应该都体会过这种场景：模型导出很顺利，板子上一跑就 OOM；或者只是把输入分辨率从 320 提到 416，内存直接翻倍。遇到这种问题，大多数时候不是算力不够，而是中间张量被无谓地“一人占一个坑”。TFLite 里负责解决这个问题的组件，就是内存规划器（MemoryPlanner）。它的角色很像推理引擎的内存管家：把不同时存在的张量分配到同一块物理内存里，只认时间不认人，用很小的开销把峰值内存压下去。这篇东西我会从推理引擎使用者的角度出发，把这个“管家”到底怎么工作、关键边界在哪、真出问题怎么排查，一次讲清楚。

很多人一开始都会问：我写推理代码从来没碰过什么内存规划器，它是怎么起作用的？其实你每次调用Interpreter::Invoke()之前，TFLite 都在后台默默做了一遍内存规划。理解了它，你以后遇到模型内存异常、自定义算子耗内存、以及部署到低内存设备时，就不会再一头雾水。

1. 内存规划器到底管什么

1.1 没有规划器的时候：张量之间各自为政

先想想最原始的方案。一个模型经过转换之后，会有一张张中间张量。每做一次卷积、激活、池化，都会产出一个或多个新的中间结果。最省事的分配方式，是给每个张量单独 malloc 一块内存，用完之后再 free。这样做逻辑很简单，代码也容易写，但代价非常直接：峰值内存等于所有张量大小的总和。

举个例子，一个只有 3 个卷积层的超小模型，中间也可能同时存在几十个张量。如果每个张量按最大尺寸预留，模型再小也可能吃掉几十 MB。这在服务端可能不算什么，但放到手机、摄像头、嵌入式设备上，就会成为致命伤。要知道端侧设备最贵的资源之一就是连续可用的 RAM，而且频繁 malloc 和 free 会带来碎片化和分配延迟，最终影响实时性。

更麻烦的是，很多模型推理是周期性的，每一次推理都会走一遍分配和释放。就算单次没有 OOM，反复分配也可能让堆里的碎片越来越多，运行几个小时后突然崩掉。这属于典型的“看起来不严重，上线后很头痛”。

1.2 规划器的核心逻辑：让不同时存在的东西住同一个房间

内存规划器的目标，不是减少单个张量的体积，而是减少同一时刻需要驻留的内存总量。打个比方：仓库里只有一个仓位，但你是早上进货、下午发货，货物不会同时出现，那这个仓位就能反复用。放到推理引擎里，张量 A 在算子 1 结束之后就再也不会被读到，而张量 B 在算子 2 才开始被创建，那么 A 和 B 就可以共用同一块内存。

判断“能不能复用”的唯一依据，是张量的生命周期有没有重叠。如果两个张量的存活区间在某个算子节点上重叠，它们就是冲突的，必须各占空间；如果没有重叠，哪怕它们看到的形状完全不同，也可以共用一块区域。

听起来简单，但真正落地时有个很现实的问题：一个真实模型有几百个张量，彼此之间的依赖关系像蜘蛛网一样复杂。人工去排根本排不过来，必须用一套算法自动计算冲突关系、排列大小、找空闲区间。这就是 TFLite 源码里MemoryPlanner接口和ArenaPlanner实现存在的意义。

2. 核心机制拆解：生命周期、冲突矩阵与贪心复用

2.1 先从“出生”和“死亡”说起

在 TFLite 的图模型里，每个算子都会有一群输入张量和输出张量。张量的生命周期可以很直观地定义成两个节点下标：

• first_node：张量第一次被某个算子读或写的节点下标。
• last_node：张量最后一次被某个算子读或写的节点下标。

只要这两个下标在同一个执行序列区间内，张量就必须“活着”。一个常见的误解是：输入张量只要刚进网络时使用一次，就可以立刻释放。实际上如果后面还有分支用到它，它的last_node会一直延伸到后期，导致它必须长时间占住内存。

拿一个最简单的链式结构为例：

输入 t0 -> 算子0 -> t1 -> 算子1 -> t2 -> 算子2 -> t3

假设执行号就是0、1、2。那么：

• t0只在算子 0 被使用，生命周期是[0, 0]。
• t1由算子 0 产生，又被算子 1 使用，生命周期是[0, 1]。
• t2由算子 1 产生，又被算子 2 使用，生命周期是[1, 2]。
• t3由算子 2 产生，是输出，生命周期是[2, 2]。

这里能看到一个有趣的点：t0和t2的生命周期不重叠，t1和t3的生命周期也不重叠。所以在规划阶段，t0的空间可以复用给t2，t1的空间可以复用给t3。

2.2 冲突矩阵：谁能和谁住一块

正式规划前，需要先生成一张冲突关系表。假设我们定几种具体张量尺寸：

如果按照朴素方案，峰值内存是四者之和：602112 + 802816 + 401408 + 200704 = 2007040字节，差不多 1.91 MB。

但根据生命周期判断冲突：

• t0和t1在节点 0 重叠，冲突。
• t0和t2不重叠，不冲突。
• t1和t2在节点 1 重叠，冲突。
• t1和t3不重叠，不冲突。
• t2和t3在节点 2 重叠，冲突。

有了冲突关系，规划器就可以开始复用。一种合理的分配结果：

• t0放在偏移 0，占 602112 字节。
• t1由于与t0冲突，只能放到偏移 602112。
• t2与t0不冲突，可以放回偏移 0 的旧区域。
• t3与t1不冲突，可以放到偏移 602112 的旧区域。

这样分配之后，峰值内存只取决于最长的连续占用组合：t1 + t2或者t2 + t3，大概在802816 + 401408 = 1204224字节左右，比朴素方案省了 40%。这还只是四个张量的玩具例子，真实模型中几百个张量压下来的收益会非常可观。

2.3 贪心分配与 ArenaPlanner 的实现思路

TFLite 默认的ArenaPlanner不是一次性解一个全局最优方程，而是用贪心策略逼近最优解。它先把所有需要规划的 RW 张量拿过来，按某个策略排序（常见是生命周期开始节点、张量大小、结束节点等），然后逐个为每个张量寻找可用内存区间。

伪代码大致是这样：

struct AllocationInfo { int tensor_index; int first_node; int last_node; size_t bytes; }; // 按一定顺序扫描所有张量 for (auto& info : sorted_allocations) { // 尝试在当前已分配的所有区间中，找一个足够大的空闲洞 bool placed = false; for (auto& candidate : arena_blocks) { if (candidate.free_size >= info.bytes && !conflicts_with_any(candidate, info)) { info.offset = candidate.offset; candidate.used = true; placed = true; break; } } if (!placed) { // 所有现有区间都放不下，就在 arena 尾部新增一块 info.offset = arena_total_size; arena_total_size += align(info.bytes); } }

这里有个很容易被忽略的细节：单个张量虽然不会和冲突张量分享同一开始偏移，但实际分配时规划器还会做对齐，比如按 16 字节或 64 字节对齐，保证后续算子可以使用 SIMD 指令。这也是为什么你算出来的理论峰值和实际arena_size永远会有一点出入。

ArenaPlanner和另一个老实现LinearAllocator相比，最大区别是线性分配器只做“从头往后排”，基本不回头复用已释放的区间；而 ArenaPlanner 会在生命周期模型上做回收。从空间利用率上看，ArenaPlanner 优势明显，所以常被选为默认实现。

3. 在真实推理引擎里如何用好它

3.1 入口在哪：默认 ArenaPlanner 是怎么挂进去的

大多数用 TFLite 的人并不会直接操作规划器实例，因为它在Interpreter::AllocateTensors()时就已经自动完成了。构建解释器时，比较常见的 C++ 流程是：

#include "tensorflow/lite/interpreter.h" #include "tensorflow/lite/interpreter_builder.h" #include "tensorflow/lite/kernels/register.h" static std::unique_ptr<tflite::Interpreter> BuildInterpreter( const tflite::FlatBufferModel& model) { tflite::ops::builtin::BuiltinOpResolver resolver; tflite::InterpreterBuilder builder(model, resolver); builder.SetNumThreads(4); // 线程数影响算子执行安排 std::unique_ptr<tflite::Interpreter> interpreter; if (builder(&interpreter) != kTfLiteOk) { return nullptr; } return interpreter; }

表面上代码里没有出现MemoryPlanner，但在AllocateTensors()内部，TFLite 会默认创建 ArenaPlanner，并为所有类型为kTfLiteArenaRw的张量分配偏移。这类张量通常就是算子之间的中间激活值，不会单独去 malloc。

如果你在部署时想确认规划器到底是否生效，可以打印每个张量data.raw指针相对 arena 基址的偏移。注意不要直接拿指针地址做业务逻辑，只用于调试：

for (int i = 0; i < interpreter->tensors_size(); ++i) { TfLiteTensor* tensor = interpreter->tensor(i); if (tensor->allocation_type == kTfLiteArenaRw) { const char* raw = tensor->data.raw; const char* base = nullptr; for (int j = 0; j < i; ++j) { if (interpreter->tensor(j)->allocation_type == kTfLiteArenaRw) { const char* p = interpreter->tensor(j)->data.raw; if (!base || p < base) base = p; } } if (base) { LOG(INFO) << "tensor " << i << " offset = " << (raw - base); } } }

如果多个张量打印出来的 offset 相同，说明它们确实被规划器复用到了同一块内存区域，这个“内存管家”在工作。

3.2 一个可复现的峰值内存估算过程

我建议每位做部署的人，在模型正式上车之前都手动估算一遍中间峰值。虽然拿工具测量内存最准，但手动估算能帮你建立直觉，避免被设备上的环境变量干扰。

还是用上一节的四张量案例。分配结果里t1占掉一块 802816 字节的区域，t2占掉一块 401408 字节的区域，而t3复用了t1的区域。于是峰值内存取t1区域和t2区域的最大同时占用，大约是 1.2 MB。如果想把峰值再降，常见路线有两类：

1. 减小单个张量的字节数。把 float32 换成 int8、int16 或者使用量化感知训练，这是最直接的手段。同样的t1，从 802816 字节直接降到 200704 字节，峰值内存会同步下降。

2. 调整生命周期结构。有些模型为了提高精度采用双向结构或频繁跳边，会把张量的last_node拉得很长，冲突变多，可用空间变少。遇到这种模型，重新设计网络结构通常比硬调 planner 更有效果。

一个我在实际优化中经常用的经验：把模型输入分辨率从 416 降到 384，看着只降了 8%，但某些中间张量尺寸会同时跨height和width两个维度缩小，再加上复用关系变化，整体峰值内存可能下降 20% 以上。所以在满足精度指标的前提下，输入分辨率往往是最划算的内存调节旋钮。

3.3 实际调优时你更该做什么

很多人会问：既然默认是 ArenaPlanner，那我能不能换一个更厉害的规划器？说实话，官方稳定 API 里并不会有让你塞自定义规划器的开关。你要么直接用默认方案，要么 fork 一份源码，在构建解释器的路径上替换 planner 实现。我不建议普通业务团队在这上面投入，因为你大概率会遇到一堆算子类型兼容问题。

更务实的调优方向沿着这个顺序来：

• 固定输入 shape，避免动态维度。kTfLiteDynamic张量常常不做 arena 复用，临时 malloc 会造成不可控的延迟和碎片。
• 调低 batch size。端侧基本用 batch=1 足够，别图省事把服务端的 batch 习惯带过来。
• 量化优先。int8 模型不仅算得快，内存占用更是直接缩成 float32 的四分之一。
• 使用模型转换时保持默认优化。某些转换选项会让所有张量都保留下来，破坏复用，仅适合调试。
• 观察自定义算子。如果你的模型里接了自定义算子，却没有正确注册 shape 推导，planner 可能无法提前拿到输出尺寸，只能退化成动态分配，所有复用都白搭。

注意：preserve_all_tensors这种调试选项只用来检查单层输出，千万别带到生产环境。开了它，内存规划器几乎等于形同虚设。

4. 常见问题与排查技巧实录

4.1 模型不大，却总是 OOM

这类问题出现时，先别急着怀疑模型太大。可能是因为你的输入张量被设成了动态，或者某些中间张量的allocation_type不是kTfLiteArenaRw。遇到 OOM 时，我会先把所有张量的allocation_type打出来看一遍：

for (int i = 0; i < interpreter->tensors_size(); ++i) { const TfLiteTensor* t = interpreter->tensor(i); printf("tensor[%d] type=%d\n", i, t->allocation_type); }

kTfLiteMmapRo是只读映射，通常对应权重；kTfLiteArenaRw是规划器管辖的读写区；kTfLiteDynamic是动态区。如果大量中间结果是kTfLiteDynamic，那你的模型大概率没有完全纳入 arena 规划。原因可能是指标输入动态、自定义算子缺 shape 信息等。

4.2 为什么两个平台上峰值内存差很多

同一个模型，在开发板 A 上峰值内存 80 MB，在开发板 B 上 100 MB。除了硬件驱动的差异，最常见原因是浮点大小和对齐策略不一致。有些平台 float 是 32 位，有些环境可能改成了 long double 编译选项；还有一些厂商的 TFLite 分支会把默认对齐改成 64 字节，导致本来可以紧密排列的小张量空出大量 padding。

遇到这种差异，不用慌。用 3.1 节的 offset 打印脚本分别导出两个平台的分配表，核对一下同 index 张量的 offset 变化，就能定位出是哪个张量被额外 pad 了。

4.3 自定义算子把内存规划器搞“失明”

这是我踩过最多的坑。自定义算子的输出张量往往没有静态 shape，注册算子时如果没有给足 shape 推导信息，解释器只能在运行时拿到实际输出大小。这时 ArenaPlanner 无法提前计算偏移，只能把该张量退回动态分配。

解决办法是在模型转换阶段就提供稳定的 shape 信息，或者在你自己的算子注册函数里实现Prepare阶段的SetTensorToReadWrite与尺寸推导。如果实在无法静态推导，那就接受这个算子带来的额外内存开销，但要注意不要让它出现在高频分支上。

4.4 多线程推理时，内存会额外翻倍吗

这里需要明确：一个Interpreter实例并非线程安全。你如果开 4 个线程并行跑 4 个推理任务，正确做法是创建 4 个独立的 Interpreter。而每个 Interpreter 都有自己的 ArenaPlanner 和独立内存池，这时峰值内存确实会按线程数翻倍。

如果你想要共享权重、减少总内存，可以考虑只共享底层模型文件映射，但中间张量缓冲没法跨 Interpreter 复用。这个限制来自图状态隔离，不是规划器偷懒。真碰到并发场景，我会优先砍线程数、复用输出缓冲，而不是试图让两个解释器共享一块中间内存。

4.5 怎么快速验证“内存管家”有没有偷工减料

除了打印 offset，还有一个很直观的办法：在Invoke()之前和之后分别记录堆内存快照，看峰值是否稳定。如果是多轮调用后内存持续增长，大概率不是 planner 的问题，而是哪里的持有引用没释放。我最常用的验证流程是：

1. 固定输入 shape，单线程推理。
2. 连续调用 100 次Invoke()。
3. 观察内存曲线是否平台。
4. 如果曲线不断攀升，先检查有没有把解释器或张量引用保存在循环外。

稳定增长的红旗信号出现得越晚，越容易和碎片问题混淆。用固定 batch=1 的小模型跑上万次循环，是最省事的压力测试。

我在实际做一个检测模型部署的时候，遇到过很极端的对比：同样的 YOLOv5s 模型，朴素实现峰值 120 MB，静态输入 + 默认 ArenaPlanner 直接压到 40 MB 左右。中间没有换任何神秘算法，就是保证张量能被规划器统一管理。所以我对内存规划器的最深体会是：它不是一个需要你手写调参的功能开关，而是一项需要你顺着它的规则来配合的基础机制。你越了解它的生命周期和复用逻辑，就越知道在模型转换和算子实现上避哪些坑。希望这篇笔记能帮你少走点弯路，下次再看到 OOM，心里能有一张更清晰的路线图。