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的空间可以复用给t2t1的空间可以复用给t3。2.2 冲突矩阵谁能和谁住一块正式规划前需要先生成一张冲突关系表。假设我们定几种具体张量尺寸张量生命周期大小float 数字节数float32t0[0, 0]150528602112t1[0, 1]200704802816t2[1, 2]100352401408t3[2, 2]50176200704如果按照朴素方案峰值内存是四者之和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 优势明显所以常被选为默认实现。特性LinearAllocatorArenaPlanner分配方式顺序累加偏移基于生命周期循环复用空间利用率较低较高分配耗时极低略高但在可接受范围适合场景简单模型/调试复杂模型/正式部署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_ptrtflite::Interpreter BuildInterpreter( const tflite::FlatBufferModel model) { tflite::ops::builtin::BuiltinOpResolver resolver; tflite::InterpreterBuilder builder(model, resolver); builder.SetNumThreads(4); // 线程数影响算子执行安排 std::unique_ptrtflite::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。如果想把峰值再降常见路线有两类减小单个张量的字节数。把 float32 换成 int8、int16 或者使用量化感知训练这是最直接的手段。同样的t1从 802816 字节直接降到 200704 字节峰值内存会同步下降。调整生命周期结构。有些模型为了提高精度采用双向结构或频繁跳边会把张量的last_node拉得很长冲突变多可用空间变少。遇到这种模型重新设计网络结构通常比硬调 planner 更有效果。一个我在实际优化中经常用的经验把模型输入分辨率从 416 降到 384看着只降了 8%但某些中间张量尺寸会同时跨height和width两个维度缩小再加上复用关系变化整体峰值内存可能下降 20% 以上。所以在满足精度指标的前提下输入分辨率往往是最划算的内存调节旋钮。3.3 实际调优时你更该做什么很多人会问既然默认是 ArenaPlanner那我能不能换一个更厉害的规划器说实话官方稳定 API 里并不会有让你塞自定义规划器的开关。你要么直接用默认方案要么 fork 一份源码在构建解释器的路径上替换 planner 实现。我不建议普通业务团队在这上面投入因为你大概率会遇到一堆算子类型兼容问题。更务实的调优方向沿着这个顺序来固定输入 shape避免动态维度。kTfLiteDynamic张量常常不做 arena 复用临时 malloc 会造成不可控的延迟和碎片。调低 batch size。端侧基本用 batch1 足够别图省事把服务端的 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 的问题而是哪里的持有引用没释放。我最常用的验证流程是固定输入 shape单线程推理。连续调用 100 次Invoke()。观察内存曲线是否平台。如果曲线不断攀升先检查有没有把解释器或张量引用保存在循环外。稳定增长的红旗信号出现得越晚越容易和碎片问题混淆。用固定 batch1 的小模型跑上万次循环是最省事的压力测试。我在实际做一个检测模型部署的时候遇到过很极端的对比同样的 YOLOv5s 模型朴素实现峰值 120 MB静态输入 默认 ArenaPlanner 直接压到 40 MB 左右。中间没有换任何神秘算法就是保证张量能被规划器统一管理。所以我对内存规划器的最深体会是它不是一个需要你手写调参的功能开关而是一项需要你顺着它的规则来配合的基础机制。你越了解它的生命周期和复用逻辑就越知道在模型转换和算子实现上避哪些坑。希望这篇笔记能帮你少走点弯路下次再看到 OOM心里能有一张更清晰的路线图。