1. 为什么推理引擎要专门配一个“内存管家”先说个我的真实经历。之前把一个视觉检测模型部署到某款低内存设备上模型文件本身只有 6.8MB量化后甚至不到 2MB我原本以为内存完全不是问题。结果第一次跑推理设备直接提示内存不足进程被杀掉。后来细查发现模型权重确实不大真正吃内存的是推理过程中产生的几百个中间张量。这里就要聊到 TFLite 内存规划器了。TFLite 作为移动端和嵌入式场景使用最广的推理引擎之一它的内存管理机制并不复杂但非常关键。官方文档和社区资料里通常会把它叫做 memory planner也就是内存规划器。通俗一点说模型执行期间有大量中间计算结果需要临时保存给这些临时数据分配空间、复用空间、回收空间的工作全部由它负责。你可以把它理解成推理引擎里的“内存管家”所有临时内存的进出都要过它这一关。这篇内容适合做移动端部署、嵌入式模型优化、或者想深入理解 TFLite 运行时的工程师。看完之后你至少能搞清楚三件事TFLite 内存规划器到底在算什么、为什么很多模型开了 arena 规划之后内存峰值能降一半、以及遇到内存规划没生效时应该从哪里开始排查。很多人会把模型占用内存直接等同于权重文件大小这是最常见的误解。权重只是加载到内存里之后变成持久缓冲区的一部分真正的峰值压力往往来自中间特征图。一个 640x640 的输入经过若干层卷积之后每一层输出的 feature map 如果全部独立分配内存会非常吓人。比如一个看起来不大的 MobileNet 系列模型中间张量全部分开存放峰值可以轻松突破 500MB这在手机上根本没法用。TFLite 内存规划器的核心工作就是通过分析张量的生命周期让互不重叠的中间结果共用同一块内存从而把这个峰值大幅度压下来。所以回到最初的场景我那 6.8MB 的模型加载之后明明只有几十个算子怎么就会把内存吃满答案正是中间张量的分配方式太粗放了。在没有合理规划的情况下每个算子输出都临时分配一块独立空间积少成多峰值自然会失控。2. “内存管家”的基础逻辑先算生命周期再安排床位2.1 张量生命周期是怎么算出来的内存规划器的第一步工作是搞清楚每个中间张量“从什么时候开始存活、到什么时候可以去世”。这里的“去世”不是说张量被销毁而是说它不再被任何后续算子引用这块空间可以拿给别的张量用。具体细节其实很直观。推理过程是一张执行图图中的每个节点对应一个算子每条边对应一个张量。某个算子的输出张量会从该算子执行结束开始存在一直存活到所有消费它的算子执行完毕。举个例子假设算子 A 的输出 tensor t1 被算子 B 和算子 C 使用那么 t1 的生命周期就是[A执行结束, C执行结束]。只要两个张量的生命周期没有重叠它们就可以复用同一块内存区域。这跟宿舍分配床位很像。同一张床上午住的是张三下午住的是李四只要两个人不会同时出现在房间里床位就没必要准备两张。内存规划器要做的就是把整个执行图上所有张量都算一遍看看哪些张量可以“错峰住同一张床”。2.2 Arena 分配一块连续的内存池TFLite 里最常见的规划方案是 arena 分配。所谓 arena就是预先申请一大块连续内存后续所有临时张量都以“在这块内存里取偏移量”的方式获得空间而不是每次独立向操作系统申请。好处有两个一是减少 malloc 调用次数二是保证所有中间结果都落在连续地址上有利于 CPU 缓存命中也方便一些底层算子做向量化。这里有个容易混淆的点。很多人以为开了 arena 就自动省内存其实 arena 只是把分配方式从“每次独立申请”变成“在一块大池子里切豆腐”。真正省内存的是规划算法。如果算法很蠢在 arena 里从头到尾顺序切块那一样会缓慢涨到高水位。TFLite 默认使用的 arena 规划器会基于张量生命周期做贪心复用所以才能真正把峰值压下来。常见的实现里有两类一类是比较简单直接的 SimpleMemoryPlanner另一类是更流行的 ArenaPlanner。两者的差别在于对执行过程的理解深度。简单规划器偏向于一次性根据执行计划算好偏移Arena 规划器则会考虑节点执行顺序、临时 buf 分配粒度等细节对复杂图更有优势。目前大多数 TFLite 版本默认用的就是 ArenaPlanner 这类实现。2.3 一个最小示例看清复用过程假设一张执行图里有这样几个张量t1算子 A 的输出后面被 B 和 C 使用t2算子 B 的输出后面被 D 使用t3算子 C 的输出后面被 D 使用t4算子 D 的输出是最终结果执行顺序是 A - B - C - D。t1 的生命周期从 A 结束一直到 C 结束t2 的生命周期从 B 结束到 D 执行前一刻t3 的生命周期从 C 结束到 D 执行前一刻。t2 和 t3 虽然都存活了一段时间但它们没有交集t2 在 C 开始时就已经没人用了所以 t3 可以复用 t2 用过的偏移量t4 出现时 t3 还没有完全放掉因此 t4 需要另找地址。如果把地址分配写出来大致是这样的张量生命周期分配的偏移量说明t1A结束 ~ C结束0x100独占t2B结束 ~ C开始前0x200独占t3C结束 ~ D开始前0x200复用 t2 的地址t4D执行中0x300新偏移这块逻辑看起来简单真正落地时还要考虑对齐、算子需要额外 scratch buffer 等情况但核心思路不会变先算出存活区间再决定谁和谁能共享同一个偏移。2.4 持久内存和临时内存要分开看推理引擎里不是所有内存都能随意复用。模型权重、常量张量、部分算子需要的固定工作区这些通常要存活整个推理生命周期属于持久内存。临时中间张量则属于临时内存可以充分复用。TFLite 内部会把这两类分开管理。持久内存基本在模型加载和 AllocateTensors 阶段就固定下来之后不会参与复用临时内存只在单次推理期间存在推理结束、再次 Invoke 之前都可以被反复改写。这也是为什么你可以在不重新 AllocateTensors 的情况下反复调用 Invoke内存占用却不会持续上涨。理解这个区分对排查内存问题特别重要。如果你看到模型加载后内存就很高那不是规划器的问题多半是权重或者常量张量太大如果你看到每次 Invoke 之后内存缓慢增长那才需要怀疑临时内存没有正确回收或者某个自定义算子在外面自己偷偷申请了内存。3. 实测从 212MB 到 76MB内存曲线的完整记录3.1 先建立基线数据做内存优化第一步永远是先测基线。我当时把一个模型跑在 Linux 的 x86 开发板上先用最原始的配置加载模型并 Invoke 一次然后在进程内部记录实际物理内存占用。代码比想象中简单本质上就是在不同节点记下当前的 RSSsize_t current_rss() { // 读取 /proc/self/statm 或调用 getrusage 获取驻留内存 // 这里省略具体实现 return 0; } std::unique_ptrtflite::Interpreter interpreter; tflite::InterpreterBuilder builder(*model, resolver); builder(interpreter); // AllocateTensors 之前的基线 size_t before_alloc current_rss(); interpreter-AllocateTensors(); // 第一次 Invoke 之后的峰值 size_t after_invoke current_rss();当时测出来的数据让我印象很深模型加载后 RSS 已经到 88MB这不奇怪因为权重加解释器本身就有基础开销但调用一次 Invoke 之后RSS 直接跳到 212MB然后就不会再回落多少。也就是说单次推理的峰值增量超过 120MB。这个增量基本都来自中间张量的独立分配。3.2 开足 arena 优化后的对比TFLite 默认其实已经开启了 arena 规划但我在早期版本里显式关闭过它用来对比开启和不开启的差距。关闭之后内存分配退化为更原始的按需分配每个中间 tensor 都有一块自己的地址空间。对比结果如下配置AllocateTensors 后 RSSInvoke 后峰值 RSS峰值额外增量关闭 arena 规划92MB212MB约 120MB开启默认 arena 规划88MB76MB约 -12MB开启默认 固定输入尺寸88MB71MB约 -17MB注意这里有个看起来反常的点AllocateTensors 之后开启规划的版本 RSS 反而从 92MB 降到了 88MB。原因是规划器在分配阶段就能识别出哪些中间张量可以共享于是不需要为每个张量单独预留一整个缓冲块共享之后预分配的总空间变小峰值自然下降。后来把输入尺寸固定又优化了大约 5MB。动态尺寸会迫使规划器按最大尺寸预留空间如果多数推理实际用的是小尺寸这部分空间就白白浪费了。3.3 用地址相等来验证复用是否生效内存规划器有没有真正干活有一个很直观的验证方法打印两个生命周期不重叠的张量地址。如果它们指向同一个地址说明规划器成功复用如果地址完全不同说明规划器没有把这两个张量重叠起来。我在程序里加过一段调试代码看起来大概是这样的// 打印某两个已知生命周期不重叠的中间张量地址 auto* t_a interpreter-tensor(tensor_a_index); auto* t_b interpreter-tensor(tensor_b_index); printf(tensor_a addr: %p\n, t_a-data.data); printf(tensor_b addr: %p\n, t_b-data.data);当时打印出来的结果里一组本应复用的张量地址完全相同说明规划器确实生效了。这个方法我现在也会用因为很多内存问题不是“规划器没开”而是“规划器规划了但某些算子让张量生命周期看起来比你想象中更长”导致复用没有发生。地址对比能够非常快地暴露这一点。实测过程中还有一个容易忽略的细节不同 delegate 的加入会影响内存曲线。如果你把一部分算子交给 GPU 或 NPU这些算子上的输入输出张量可能不再进入 CPU 规划器的 arena而由 delegate 自己管理。此时你看到的总体内存增量实际上是 CPU 临时内存和 delegate 内部 bufffer 加在一起的结果。做对比时最好先把 delegate 关掉做一次基线再逐步打开否则很难判断优化到底发生在哪一层。4. 规划器“失灵”的三个典型场景4.1 动态 shape 让规划器只能按最大尺寸做事很多模型在部署时输入尺寸不定比如目标检测模型允许不同分辨率输入。TFLite 对此有 ResizeInputTensor 接口。问题在于每次 resize 之后所有中间张量的 shape 也会跟着变规划器必须重新计算 offset。如果模型里某几个中间张量在某个固定 shape 下很小但在另一个 shape 下很大规划器为了安全只能按最大情况预留空间。我踩过的坑是模型默认支持 1920x1080 输入但我实际部署只需要 640x640。因为我没有在加载模型后显式把输入 resize 到 640x640规划器一直按 1920x1080 的最坏情况来分配 arena。结果就是模型明明是轻量的内存峰值却比预期高了好几倍。解决办法很简单AllocateTensors 之前先调用 ResizeInputTensor 把输入固定成实际使用尺寸。如果产品端确实需要多种输入尺寸最好在内存规划时预留一个统一的最大缓冲或者在运行时按不同尺寸分别维护一份 interpreter避免反复 resize 带来的重新规划开销。4.2 多子图和条件分支导致生命周期无法全局计算TFLite 模型不一定是一张简单的图。遇到控制流算子比如 while 循环、条件分支或者一些算子被拆到多个 subgraph 中执行时内存规划器默认只能基于主图的静态执行计划做分析。子图里的张量生命周期很难和主图里的张量做完美交叉复用。这种情况在实际中很常见尤其是从 TensorFlow 转换过来的模型经常带有控制流逻辑。表现出来就是整体看模型不大内存却一直在高位徘徊。我当时的排查思路是先尽量把模型里的控制流移除让执行路径变成单一静态路径。比如把 while 循环在转换前展开成固定次数的算子序列或者把条件分支里不会执行到的那一侧直接裁剪掉。这样做之后规划器能看到的执行计划更准确内存复用率也会明显提升。4.3 自定义算子外面的“看不见”内存TFLite 内存规划器能管理的只是框架层面能看到的张量。如果你在模型中加入了自定义算子这个算子在执行时自己又 malloc 了一块很大的临时 buffer那这块内存完全处于规划器的视野之外。规划器不会帮你复用也不会主动释放全部要看算子作者自己的小程序写得好不好。我在一个后处理自定义节点上遇到过一次问题。那个节点内部对整张 feature map 做了复制用来做 NMS 前的排序操作。表面上那只是一次临时的算法开销但每次推理都临时分配一整块内存积少成多峰值直接顶到设备上限。后来我把这块临时内存改成在算子第一次执行时就预分配好之后反复复用问题立刻消失了。因此排查 TFLite 内存问题时如果框架层 arena 已经没有明显优化空间下一个要查的地方一定是自定义算子内部的分配逻辑。5. 当默认规划器无法满足要求时还能做什么5.1 优先改图而不是强行改规划器很多人在内存吃紧时第一时间想的是怎么写一个更聪明的内存规划器。说实话这条路能走通但不划算。TFLite 默认的规划算法在绝大多数静态图上已经接近最优解因为张量生命周期分配问题本质上可以转换为区间图着色问题。区间图着色是多项式可解的贪心算法就能得到最优解。也就是说只要执行计划是静态的默认规划器算出来的结果基本就是你能达到的最优复用方案。所以如果图本身是静态的、张量生命周期一目了然内存依然不够问题多半出在图的结构上而不是规划器的能力上。这时候更有效的做法是回头审视模型能不能少几个分支能不能把某些大 feature map 在后处理时尽早释放能不能用一些算子合并或裁剪技术减少中间张量的数量我在多个项目里发现改图往往比改规划器收益大得多也稳定得多。5.2 自定义内存规划器的思路与代价确实有些场景默认规划器不适合。比如你的模型里有一个算子需要非常大的对齐缓冲区或者某个 intermediate tensor 的生存期极短、但尺寸极大默认贪心策略没有按照你的业务约束做特殊处理。这时候可以考虑写一个自定义规划器。TFLite 源码里给规划器留了扩展接口。想自己实现核心要做的其实也就两件事第一明确每个 tensor 的生命周期第二在生命周期不冲突的前提下为各 tensor 分配尽量紧凑的内存偏移量。听起来容易实际上要考虑对齐、算子执行顺序中 astaging buffer 的临时占用、以及子图间往返执行等因素。实现一个能跑通的 demo 可能只需要几个小时但要覆盖各种边界情况并保证不出现数据互相覆盖需要大量的迭代验证。我的建议是先读一读 arena_planner.h 的实现模仿它的接口写一个最小版本在测试集上验证一段时间。如果只是为了让某一个特定算子省点内存完全可以通过修改自定义算子内部实现来达到目的没必要动规划器这层。5.3 大缓冲区延迟释放用小技巧解决大问题另外还有一种比较隐蔽的场景模型里有几个中间张量超级大但在推理最后阶段已经不被使用不过框架层面它们的生命周期绑定了整个 subgraph 的执行范围一直到最后也没法释放。对于这类中间结果可以在图结构上手动拆出独立子图或者通过自定义算子在不需要的时候提前释放内存。虽然这种做法稍微有点 hack但在生产环境中确实有效。我当时处理过一个模型最后阶段拿到检测结果后有一个 20MB 的中间 feature map 其实完全没用了但它属于整图最后一个节点之前的部分规划器认为它必须活到那个节点执行完。后来我在它最后一次被消费之后插入了一个自定义算子这个算子的工作就是把这个张量的缓冲区重新标记为空闲并在后续推理中复用给其他中间结果。就这么一个小改动峰值内存又降了十几兆。6. 落地前我给自己列的内存经验清单做完上面这些优化之后我把整个排查思路收敛成了一份简单的清单。每次遇到 TFLite 内存问题我都会按这个顺序过一遍先把输入尺寸固定成实际部署尺寸不做多余预留。确认默认 arena 规划已经开启并且 AllocateTensors 只调用一次。用地址对比法抽查几个生命周期不重叠的张量确认复用确实发生。关掉所有 delegate测一次 CPU 基线再逐个打开 delegate看内存变化发生在哪一层。检查自定义算子内部有没有每次 Invoke 都重复申请的临时内存。如果模型带控制流优先考虑裁剪或展开让执行路径尽量静态化。考虑模型结构层面是否有多余的中间张量可以合并或提前释放。这套流程并不是什么高深理论更多是实践经验。第一次做内存优化的人最容易犯的错误就是只盯着模型文件大小或者只盯着框架 api 的内存统计数字忽略了内存规划器的实际行为。真正落地时你需要的不是把所有内存都压到极限而是让峰值稳定可控、不出现偶发性的 OOM。最后再分享一个个人体会调内存规划器这种事很多时候你并不会真的去写一个新规划器而是要做一个不讨喜但正确的事——回头改模型结构。把一个看似“已经训好的模型”再拆开、裁剪、重置中间张量听上去是在退步但实际上推理引擎会给出非常实际的回报。内存降下来的同时因为缓存局部性变好推理延迟偶尔也会跟着下降。所以别把“内存管家”当单纯的分配器看它其实是帮你检验模型是否适合端侧运行的那么一位严管家。