TFLite内存规划器：端侧推理中不可忽视的内存管家

做端侧推理这一年多我越来越觉得 TFLite 的“内存规划器”才是推理引擎里那个最容易被忽视的“内存管家”。很多人把 TFLite 当成黑盒加载模型、AllocateTensors、然后舒舒服服调 Invoke但很少有人会问一句——这几十 MB 的中间激活值到底住哪谁给它们办的“入住手续”如果模型跑起来总是内存暴涨、频繁被系统杀掉或是在低内存设备上调来调去都没效果那问题大概率不在模型本身而在于你没有理解这个“管家”的工作方式。这篇文章我就以实际调优踩坑的经历为主线把 TFLite 内存规划器从职责、算法、观测手段到调优方向完整拆开讲一遍。适合正在做模型部署、搞端侧推理、或者被 native heap 涨跌折腾到怀疑人生的开发者。1. 先聊一个问题为什么 TFLite 不用 malloc 一把梭1.1 那些年我见过的“同模型不同内存”现象我最早注意到内存规划器这个东西是在调一个 OCR 检测模型的时候。同一个模型、同一份权重两个同事写的推理代码一个在低端机上跑得稳另一个在多 4GB 内存的测试机上反而 OOM。我当时觉得很奇怪模型文件才 12MB激活值再怎么算也不至于让 6GB 内存的手机扛不住吧后来仔细对比才发现问题出在对“中间张量”的处理上。一个模型基本上都是几十上百个算子串起来每个算子都会产生输出张量这些中间张量如果每个都当成独立内存去申请峰值内存就会是“所有算子输出之和”这跟“同一时间真正存活的张量总大小”完全是两码事。TFLite 之所以能在端侧跑起来核心原因就是它有专门的内存规划器能在模型加载后做一次完整的“内存排班”让不同时间存活的内存互相复用。那次之后我养成了一个习惯任何 TFLite 模型跑起来之后第一件事不是看推理耗时而是看内存峰值。这个峰值如果不正常后面所有性能优化都是白做。1.2 权重与激活值两笔账必须分开算聊内存规划器之前先把概念理清楚。一个模型在推理时占的内存大致分两类权重常量张量卷积核、偏置、BN 里的 mean/variance 这些。它们是不可变的模型加载后整块存在内存里重复用。激活值中间张量每一层算出来的中间结果。它们论“辈分”存在——前一个算子算完给后一个算子用用完之后就没用了。模型文件再大如果所有权重都被只读映射或直接保存在内存里它占的是“静态账”。真正容易把内存打爆的是“动态账”也就是激活值。内存规划器主要管的也是激活值这一块。打个比方权重就像仓库里囤的货激活值就像快递分拣线上短暂停留的包裹。货可以一直堆在那包裹却得在流水线上流转你不做管理积压的包裹就会把整个仓库塞满。1.3 通用分配器做不到的“预知未来”那为什么不干脆让系统的 malloc / new 自动管理系统分配器当然也能分配内存但它有一个致命弱点它看不到计算图的全局信息不知道某个张量到底什么时候死。举一个极端点的情况。一个模型的输入是 1×192×192×32 的 tensor假设是 float32大小约 4.7MB。这个 tensor 生成之后有 6 个分支算子都要读它。如果按“谁用谁分配”的思路每个分支算子内部可能都会为它的输出单独搞一块buffer于是同一时刻活跃内存翻好几倍。但内存规划器就不一样。它在真正推理之前已经预读了整张计算图的执行顺序能给每个张量算出存活区间开始于哪个算子、终结于哪个算子、哪些张量永远不会同时存在。只要两个张量的存活区间不重叠它们就可以共用同一块内存。这种“预知未来”的能力是通用内存分配器不具备的也是 TFLite 能在几百 KB 内存上跑模型的基础。2. 内存规划器到底管什么我给每个中间张量排了张时间表2.1 从“AllocateTensors()”开始看内存从哪来TFLite 的调用流程大家很熟加载模型、配置解释器、调 AllocateTensors、然后 Invoke。很多教程把 AllocateTensors 简单解释成“分配模型所需内存”但实际上这一步做的事非常关键它是在执行内存规划器的“排班表”。第一次调用AllocateTensors()时解释器内部会做这几件事拿到执行计划ExecutionPlan也就是算子按照拓扑序排好的执行顺序。根据每个算子的输入输出把每个张量的“出生点”和“最后使用点”都标记出来。把所有非常量张量放进同一个 arena竞技场/内存池里统一规划偏移量。一次性向系统申请一整块大内存然后把这整块内存按偏移量划分给各个张量使用。这里最反直觉的一点是AllocateTensors()不是在为每个张量单独申请内存而是向系统要了一次“打包批发”。所有人都住在同一栋楼里只是楼层不同。如果你在 Android 上做内存分析会发现 AllocateTensors 之后 native heap 出现一次很大的凸起之后不管怎么 Invoke内存都不会再有大波动——因为后面全在内部划桌子不找系统要新地了。2.2 活跃区间决定一个张量何时“该死”内存规划器做排班依赖一个概念活跃区间live range。一个张量的活跃区间大致是这样的出生第一个写它的算子的执行完成时刻。存活它作为后续若干算子的输入而存在的整个时间段。死亡最后一个读它的算子执行完成时刻。只有在“出生之后、死亡之前”这段时间里张量才真正占着内存。它一旦死了它占的这块地就可以被之后出生的张量继承。TFLite 源码里做这一件事的地方叫 GraphInfo 或者模型上下文分析它从执行计划反推出每个 tensor 的 first_create_node 和 last_use_node。这个过程跟编译器里的活跃变量分析非常像——写过 LLVM 或做过寄存器分配的人看这段代码会有强烈的既视感。我举个例子你就明白了。假设有这么一段算子序列Conv_1 - Conv_2 - MaxPool_1 - Conv_3 - MaxPool_2 - FCConv_1 的输出如果有三个算子要用它就一直活着MaxPool_1 的输出只有 Conv_3 要用那么 Conv_3 算完那一刻它就死了它所占的存储就能让给后面的张量。推理引擎最怕的是模型里出现一种叫“共享依赖”的结构——一个 tensor 被十个算子依赖那这十个算子没跑完这块内存就永远不能挪作他用。这也是为什么很多工程师看到自己模型“分支多、树状结构复杂”时内存峰值会明显高于链式模型。分支多意味着同时存活的大张量多复用空间小。2.3 arena提前包下整块场地再内部划线讲清楚活跃区间之后再来看看 arena 这个概念。TFLite 的 arena 本质上就是一块连续的大内存。你可以把它理解成你办一场演唱会与其给每个观众单独开一个房价去订酒店不如直接包下一整栋楼然后由主办方自己去分房。省去了大量“反复找中介谈判”的开销系统调用也避免了散住导致的各种浪费内存碎片。arena 内部的所有偏移量全部由规划器事先算好。每个张量在 arena 里占哪一段、跟谁重叠彼此互不干涉。由于所有中间张量都在同一块 arena 里只要峰值算得准整个推理过程就不会再发生额外的系统级内存分配。这在嵌入式或低内存设备上尤其重要。很多实时场景下一次 malloc 可能都会引入不可控的耗时抖动而 arena “赛前规划、赛中无分配”的特点天然适合对时延敏感的应用。3. 源码视角ArenaPlanner、SimpleMemoryArena 与 GreedyMemoryPlanner 的分工3.1 首先是 GraphInfo规划器的一手数据来源如果你去翻 TFLite 源码会发现和内存规划相关的类主要这么几个ArenaPlanner、SimpleMemoryArena、GreedyMemoryPlanner以及一个给规划器喂数据的GraphInfo接口。GraphInfo做的事情非常朴素它向外暴露模型里有哪些节点、每个节点的输入输出 tensor 索引、每个 tensor 是否可写等等。ArenaPlanner 拿到这些信息后先做一轮“谁是死的谁是活的”分析把每个 tensor 的生命周期整理成一张表。这张表就是后面所有内存分配决策的依据。有个细节值得提一下GraphInfo 不是只给 tensor 的 shape它还会告诉规划器这个 tensor 是不是属于某个 subgraph。TFLite 支持多 subgraph 结构在控制流模型里这一点很重要。每个 subgraph 是独立的 arena 空间parent graph 和 child graph 之间的桥接只会走固定几个输入输出槽。3.2 ArenaPlanner 的簿记逻辑ArenaPlanner 的簿记逻辑可以理解成一张 Excel 表。它给每个 tensor 维护几列信息tensor 在计算图中的索引tensor 需不需要持久化比如输出 tensor 必须等到下一轮才能释放tensor 的字节数tensor 被分配到的 arena 起始偏移量分配的时候它会从已死亡张量留下的空洞里挑一块足够大的把新 tensor 塞进去。如果没有足够大的空洞就往后追加一块新区间。所以 arena 的整体布局并不是“按张量索引排排坐”而是完全由生命周期和大小驱动看起来非常乱但内存密度极高。一个容易误解的点是两个 tensor 复用同一块内存并不要求它们大小相等。只要后一个张量比前一个张量小或者大小恰好能塞进空洞就可以。只有所有空洞都装不下时才会往 arena 尾部扩张。所以一个模型里几 KB 的小张量和几 MB 的大张量完全可以交错挤出空间。3.3 贪心如何摆平大小不一的张量真正的底层内存规划算法叫GreedyMemoryPlanner。从名字就能看出来它不是全局暴力搜索最优解而是用贪心策略快速构建一个足够好的方案。为什么用贪心因为在一个含几百个节点、上千个张量的大模型里如果追求全局最优时间复杂度会高到无法接受。工程师真正想要的是在几十毫秒内完成规划同时让峰值内存接近最优。贪心策略配合活跃区间分析实践中已经能拿到非常接近理论上限的复用率。打个比方你把一堆形状不同的行李箱放后备箱追求绝对最优摆放需要穷举普通人就是先放大的再把小的往缝隙里塞。GreedyMemoryPlanner 就是这个“先放大件、再填缝隙”的思路而 SimpleMemoryArena 就是那个后备箱。在 TFLite Micro 里这套思路也一样。MicroAllocator 用自己的 micro planner 做同样的生命周期分析和内存复用只不过它面对的往往不是几百 MB 的内存而是几十 KB 的 RAM。原理同源工程实现上更省。4. 把账本拉出来看看实测 TFLite 内存复用4.1 最朴素的观测方法比较 data.raw 指针看再多源码不如亲手把账本拉出来。TFLite 里所有张量都对应一个TfLiteTensor结构里面有个data.raw指针指向这个张量实际使用的内存地址。如果两个张量被规划器复用了同一块 arena 内存那它们的data.raw就应该相同。我常用这样一段小代码来快速观察#include cstdio #include memory #include tensorflow/lite/interpreter.h #include tensorflow/lite/kernels/register.h #include tensorflow/lite/model_builder.h int main() { auto model tflite::FlatBufferModel::BuildFromFile(model.tflite); tflite::ops::builtin::BuiltinOpResolver resolver; tflite::InterpreterBuilder builder(*model, resolver); std::unique_ptrtflite::Interpreter interpreter; builder(interpreter); interpreter-AllocateTensors(); printf(tensors_size%d\n, interpreter-tensors_size()); for (int i 0; i interpreter-tensors_size(); i) { auto* t interpreter-tensor(i); if (t nullptr) continue; printf(idx%-4d bytes%-8zu ptr%p\n, i, t-bytes, static_castvoid*(t-data.raw)); } return 0; }注意头文件路径以你用的 TFLite 版本为准较新版本用model_builder.h老一些的版本是model.h。跑完输出之后把所有ptr相同的张量归成一组你会发现一个很有意思的现象内存里真正独立的存储块数量远比张量数量少得多。这就是内存规划器在背后干活的最好证据。我自己通常会把输出整理成一张小表下面是某种典型中型模型可能出现的情况Tensor 索引字节数data.raw 指针备注318432000x7f80c10000Conv 输出818432000x7f80c10000与 3 复用同一块内存119216000x7f80c0e000Pool 输出152304000x7f80c0e000与 11 复用1245760x7f8383b000常量权重落在模型映射区如果是 float32 模型bytes算出来基本吻合C*H*W*4如果是 int8 量化模型基本吻合C*H*W*1。看到大量复用小组的时候你就可以拍胸脯说这个模型的内存规划是生效的。4.2 从一组例子里理解复用的边界光看指针还不够你得能解释“为什么这个复用成立、那个不成立”。举一个我实际遇过的例子。某个语义分割模型里输入是 1×512×512×3第一层 conv 输出 1×512×512×16大约 16.7MBfloat32。它后面接了 4 个分支每个分支都会消费这个输出。也就是说从第一层结束到第四个分支结束这 16.7MB 一直活着。然后是另一个超大张量倒数第二层输出的 1×512×512×2。它出生在最后一个分支里那个时候前面 4 个分支的中间张量基本都已死光。所以规划器就把这两块内存安排到了同一个位置。虽然两者类型长得差很远一个 16 通道一个 2 通道但只要生命周期不重叠一切好说。假设模型作者偷懒在一个分支里没有及时把输出“用掉”再进入下一个分支那最后一层的大张量就要和前面的大激活值共存峰值内存直接翻倍。这种问题靠肉眼很难在 tensorboard 图里看出来但你把这组 data.raw 拉出来一眼就知道哪块内存没被复用。所以做内存排查我建议别急着上 profiler先把每个 tensor 的占用和指针 dump 一遍把所有复用关系画清楚。往往问题就出在最不起眼的几个“存量大、存活时间长”的张量身上。4.3 如何用 SetCustomAllocationForTensor 接管输入输出内存还有一些场景你想自己直接掌控某块内存。TFLite 提供了SetCustomAllocationForTensor允许你给指定 tensor 手动指定一块内存而不是让规划器在 arena 里安排。最典型的使用场景是输入图像已经从摄像头 DMA 拿到一块内存里了你不想再拷贝一次到 TFLite 的输入张量缓冲区直接把这块 DMA 内存交给输入张量省一次拷贝。给输出张量指定内存也一样数据可以直接写到预先分配好的 buffer 里之后直接进渲染管线。用法上大致是先构造一个TfLiteCustomAllocation把data和bytes填好再传给解释器。需要注意自定义分配的内存首地址要对齐尤其是后面要接 GPU delegate 或者 DSP 时对齐要求更严格另外一旦你给某个张量设置了自定义分配TFLite 就不再参与这块内存的生命周期管理释放必须由你自己来。也就是说“管家”愿意把某些房间的管理权交给你但一旦出了事责任也全在你。5. 规划器最怕的三件事动态形状、委托和多线程工作区5.1 动态形状规划退化成“最坏情况预留”内存规划器的一切精确计算都建立在“张量形状已知”这个前提上。你在训练脚本里固定了输入尺寸模型转换之后就带着这个尺寸固定下来规划器才能精确算字节数、做复用。但如果模型里带有动态形状比如输入尺寸可变或者某个循环次数取决于运行时数据规划器就没办法在 AllocateTensors 的时候知道最终大小它只能退化成另一种策略每个动态张量都按最坏情况预留。最坏情况是什么对很多视觉模型来说输入从 192 变成 640激活值直接涨十倍都不止。更麻烦的是如果你用ResizeInputTensor反复改变输入尺寸TFLite 每次都会重新规划 arena。重新规划本身不慢但它意味着之前记忆化的内存布局全部作废也许还会触发一次更大的底层分配。所以我会给同行一个非常朴素的建议能定死输入尺寸就定死实在要动态尽量限制在有限的几个档位而不是任意尺寸。固定档位可以让规划器在每档内部依然做到精确复用任意尺寸则几乎逼着它时时做最坏打算。5.2 delegate 接管后arena 外又开了一本账TFLite 的加速很多时候靠 delegate比如 Android 上的 NNAPI、GPU delegate、XNNPACK 等等。delegate 通过ModifyGraphWithDelegate把部分算子“外包”给硬件或更高效的后端。外面看着是同一个模型内部内存管理其实已经变了。delegate 有两种常见做法把算子的中间 tensor 直接留在自己的私有内存池里完全绕过 TFLite arena只把输入输出 tensor 与 TFLite 做桥接中间过程全部黑盒。这就导致一个现象你在data.raw里仍然能看到模型中间表示的逻辑 tensor但它们的指针可能指向 delegate 内部缓冲和普通张量之间不再有清晰的复用关系。比如 NNAPI delegate 运行时硬件驱动可能自己管理内存TFLite planner 对它控制的算子区间是“看不见”的。所以如果你在测试里发现“加了 delegate 后 native heap 反而多了一块不明内存”别急着怀疑内存泄漏。先确认一下是不是 delegate 在工作区间内建了私有池。多数情况下这是正常甚至是必要的——GPU 和 DSP 需要连续、对齐、甚至带 cache hint 的专用 bufferarena 里的普通内存根本没法直接用。5.3 多线程工作区账单上没有的隐性支出内存规划器管的是 tensor 的内存但它管不了算子内部的“工作区”。这句话我在实际调优中吃了不少亏。有些算子尤其是深度的卷积、矩阵乘它们在 C 实现里会为多线程分块计算开辟临时 buffer。这个 buffer 的生命周期完全存在于单次算子调用内部不经过 arena也不体现在 tensor-data.raw 上。如果你把set_num_threads从 1 调到 4会发现内存峰值可能跟着涨——因为多线程需要更多分块 buffer、可能需要做 padding 对齐每个线程可能还要有自己的累加 buffer。这类内存不在内存规划器“账本”上但它是真实存在的支出。排查内存问题时如果 tensor 复用好得无可挑剔但内存还是高先看一眼是不是线程数太多或者某个自定义算子在内部做了大量临时分配。对应地在内存极度紧张的设备上不要迷信“线程越多越快”。线程数加倍内存可能多出一块推理速度却不一定有正收益。这种多线程工作区通常受益于线程池首次调用后保持不变所以如果你要做内存峰值压测一定要跑足够充分的 warm-up否则看到的峰值是偏小的。6. 把峰值内存压下去几个经过验证的调优方向6.1 结构优先减少同一时刻还活着的张量前面讲了那么多机制落到实践上第一个调优方向永远是结构性的减少同时存活的张量数量。同一个模型理论上等价的不同结构内存峰值可能相差三四倍。典型例子是一个超大 tensor 被多个分支共享能不能改成“先用一块再算另一块”分支里的输出是不是都必要有些只是为了监控可视化而保留的中间输出部署时完全可以剪掉。能不能把“串行的小算子组”合并成一个大算子这些改动不是 TFLite 层面能替你完成的属于模型结构设计。但你可以通过 dump 出来的复用关系精准找到“活得又久、占得又大”的张量再回头和算法同学商量怎么改。另一个实际可行的点是执行顺序。TFLite 执行计划通常按拓扑序固定下来但如果你有多个互不依赖的分支它们的先后顺序会影响峰值。两个大分支先算 A 再算 B和先算 B 再算 A峰值不同。虽然现版本 TFLite 不直接暴露“手动重排执行计划”的稳定 API但了解这一点能帮你在分析时少绕弯——当复用关系不理想时多半是结构问题而不是规划器算法问题。6.2 量化与融合让同样一块内存装更多结果量化大概是性价比最高的内存削减手段。转换模型时加个量化float32 变成 int8单位元素从 4 字节变 1 字节激活值和权重同时缩到原来的四分之一。一块本来只能放一个大激活 tensor 的 arena现在能放四个。内存压力瞬间大减。另一个和量化配合得很好的手段是算子融合。TFLite 转换器以及一些自定义 pass 会把 Conv BatchNorm Relu 这种连招合成一个算子。从内存规划器视角看这意味着原来那个用来传递 Conv 输出的中间张量直接消失了它不用再活到后续算子都读完。我见过实测结果一个原本峰值约 900MB 的浮点分割模型做完整 int8 量化加算子融合之后峰值降到 180MB 左右。这里大头是量化但融合省下来的那部分往往是纯白捡的既不加精度损失还能减少耗时。唯一要留意的是量化模型的精度在数值敏感场景下可能不达标。建议先在测试集上对比输出误差如果误差兜不住可以先试“混合量化”或“量化感知训练”别一棒子打死整数量化方案。6.3 两块硬核配置subgraph 隔离与自定义 allocator最后讲两个进阶玩法。第一个是 subgraph 隔离。TFLite 支持多个 subgraph控制流逻辑或者某些相对独立的子模块可以拆出来它们各自拥有独立的 arena。这种隔离在某些场景下会让总内存略高因为不能跨 subgraph 复用但优势是某个 subgraph 内部做动态形状调整时不会把整个模型的顶层内存布局都搅乱。对复杂动态模型来说这种局部隔离反而是减少总峰值的手段。第二个是接入自定义 allocator。如果你用的是比较新的 TFLite 版本或者基于它魔改的运行时可以通过修改底层分配策略把 arena 的整块申请从malloc换成你自定义的分配器比如预分配静态内存池、匿名共享内存或者带对齐要求的硬件内存。在 Android 上有人会把 arena 指向ashmem这样既能控制内存上限又方便与系统共享。这个属于较深度定制适合你已经把前面几步吃透之后再去碰。总的来说我现在的习惯是模型到手先跑一遍 data.raw dump把复用关系打印出来再决定要不要动结构、要不要上量化、要不要加 delegate。这一步几乎没什么成本但能省下后面好几天的跟内存搏斗时间。内存规划器这个“管家”平时不吭声账倒是记得清清楚楚关键是你得知道去哪看账本。