第一次把训练好的模型搬上真机时我盯着内存监控面板愣了很久模型文件明明只有 12MB进程峰值却跑到 63MB。多出来的内存既不是权重也跟输入图片没太大关系大部分被推理引擎在执行过程中产生的中间张量吃掉了。后来我把 TFLite 的运行时源码翻了一遍才算搞明白这套藏在引擎底层的“内存管家”——内存规划器Memory Planner到底在干什么。这篇文章想把自己的理解讲清楚TFLite 是怎么给中间张量安排内存的、它依据什么规则复用缓冲区以及如果你要自己写一个推理引擎这些思想如何直接照搬。1. 模型只有 12MB推理峰值却要 60MB中间张量才是内存大头1.1 推理引擎的内存账单由两部分组成任何推理引擎跑起来内存占用都能拆成两块静态部分和动态部分。静态部分包括模型权重、bias、量化缩放因子这些常量它们从模型文件加载后基本不再变化动态部分则是推理过程中计算出来的激活值、临时缓冲区、算子工作区这些数据随每层计算不断产生、消费、释。TFLite 对权重这类常量做了一个很聪明的处理能 mmap 就直接 mmap 进进程地址空间不需要一次性把整个模型读进堆里再拷贝一份。所以开篇那个 12MB 的模型权重部分在内存里可能连 12MB 都不占而且这部分是只读的永远不会被回收。真正的内存大头在中间张量。每一层算子的输出对下一层来说就是输入。一个卷积跑完一份 activation map 就产生了这份数据必须要留在内存里等下一个算子把它读走。如果你一层层算下去所有中间结果都不回收那累计数量非常吓人。TFLite 的会话级峰值内存很大一部分就是被这些中间张量撑起来的。1.2 一次 Conv 就能吃掉多少内存一笔实际开销账内存大小公式其实很简单中间张量的大小一般是batch × height × width × channels × 每个元素字节数。拿一个常见的 224×224 图像分类模型来算输入是 1×224×224×3 的 RGB 图float32 下大约是 150KB很小。但第一个卷积核如果是 64 通道输出就变成 112×112×64这里就是大概 3.2MB。仅仅经过一层中间张量就放大了 20 倍以上。更麻烦的是深层网络。MobileNet 这类轻量网络虽然参数少但中间特征图的体积并不小几十层叠加所有中间张量总和动辄 60MB 到 100MB。如果是检测模型neck 部分还会出现 80×80×256 这种尺寸的特征图一张就是 6.5MB几个尺度的特征叠加起来压力更大。把 batch 调到 4这些数字全部翻四倍。所以很多做端侧部署的人第一反应都是把 batch 固定成 1分辨率能降就降这本质上就是在跟中间张量内存做斗争。1.3 朴素分配为什么在移动端根本活不下去有人会问我每个算子算完就把输出 malloc 一块内存用完了再 free不也行吗理论上行但实际移动端场景下几乎没法用。问题在于malloc/free本身是有成本的。极端情况下一个 100 层模型就会产生上百次内存分配和释放每个算子都经历一次「找堆、锁分配器、更新空闲链表」的流程推理延迟和功耗都会被拖起来。更麻烦的是碎片化——不断有大有小、有快有慢的分配释放内存堆会变得越来越碎明明总量够却分配不出一块连续的大内存。嵌入式 Linux 和 Android 又没有 swap 空间一旦 OOM进程直接被系统杀掉没有回旋余地。所以推理引擎需要一个比“随用随 malloc”更聪明的角色提前把整个执行图看清楚算出哪些内存可以复用、哪块地址给哪个张量用一次规划整个生命周期稳定运行。这就是内存规划器存在的价值。2. 生命周期冲突图为什么两个“错峰”张量可以共用一块缓冲区2.1 判断能不能共享只需要一个问题它们是否同时活着内存规划的核心不是看你总共需要多少内存而是看同一时刻最多有多少张量活着。“活着”的定义很直接一个张量从被某个算子生产出来到最后一次被某个算子读取这段区间内它就是活着的。活着期间这块数据必须待在内存里不能被覆盖一旦它的最后一个消费者跑完它就可以被视为“死亡”物理内存立刻释放。我们常说的“复用”本质就是酒店前台给两拨错峰入住的客人安排同一个房间。前一个客人中午退房后一个客人下午入住只要他们不同时待在房间里房间就可以重复用。张量共享缓冲区也是同一个逻辑如果张量 A 的存活区间和张量 B 完全不重叠那么 A 释放之后B 可以占用 A 原来的内存地址。2.2 从执行顺序反推生命周期出生点和死亡点要让规划器知道两个张量是否错峰先得知道算子的执行顺序。推理图是有依赖关系的 DAG把节点按拓扑序排好就得到一条确定的执行链。有了执行顺序每个张量就能标出两个关键点出生点——哪个节点把它生产出来死亡点——哪个节点最后一次读取它。举个例子一个简单的算子链op1 - op2 - op3 - op4假设 op1 输出 tensor Bop2 是 B 的唯一消费者op2 输出 tensor Cop3 消费 Cop3 输出 tensor Dop4 消费 D。按照执行步数来看B 的存活区间是 [1,2]C 是 [2,3]D 是 [3,4]。B 和 C 在 op2 结束前后交接实际上它们是“前后脚”的关系内存完全可以交替复用。这里有个细节很多人会忽略一个节点在执行过程中它的输入和输出是同时活着的。所以张量的出生时刻应当是“生产者节点完成之后”不是“生产者节点开始之前”。用半开区间去描述才能避免把同一节点内的输入输出误判成可共享。真实框架里规划器会按节点完成时点切分生命周期比上面的整数区间描述更精细但本质逻辑是一样的。2.3 区间图着色把“内存块”想象成教室现在问题变成了给一批带有存活区间的张量找最少的内存块让每个块在任意时刻最多只被一个张量占用。这在算法上就是区间图着色问题。我用教室来类比。每个张量是一节课上课时间是它的存活区间教室是内存块。两个课程时间不重叠就能共用同一间教室重叠了就必须分开。最少需要多少间教室就等于最大同时上课的课程数。把内存换成教室记忆负担小很多而且完全等价。对区间图来做贪心着色按开始时间排序、依次分配现有空教室结果是相对优秀的。这也是为什么很多推理引擎不整花活老老实实按“存活区间 引用计数”扫一遍就能拿到非常接近最优的结果。2.4 TFLite 里的实现方向引用计数与空闲块池落到工程实现上常见做法是引用计数。规划器按节点顺序扫描每扫到一个节点就把该节点输入张量的引用计数减一等到引用计数归零说明这个张量已经没有消费者了立刻把它占的内存块放进空闲池。处理输出张量时先看空闲池里有没有够大的块有就直接复用没有再扩展内存池。这个方案实现的复杂度不高但效果非常好。它不需要一口气把所有中间张量的大小都预留出来而是“随用随腾”自然就把峰值压到了最大重叠量附近。TFLite 的 ArenaPlanner 就是这么干的下面拆开看细节。3. ArenaPlanner 拆解TFLite 是具体怎么当这个“内存管家”的3.1 每种张量有各自的“物业类型”TFLite 里每个张量都有一个allocation_type可以理解为它的“物业类型”。类型不同内存管辖方式完全不同。类型生命周期是否参与 Arena 复用典型用途kTfLiteMmapRo整个进程生命周期不参与权重、常量直接映射模型文件kTfLiteArenaRw单次推理内动态产生和释放参与绝大多数中间激活张量kTfLiteArenaRwPersistent跨多次推理/跨子图存活不参与普通复用模型状态、循环变量、跨子图共享缓冲kTfLiteDynamic运行时可变不参与静态规划形状动态改变的输出、NMS 后处理输出这里最值得注意的其实是kTfLiteArenaRwPersistent。它虽然也分配在 Arena 里但不会像普通中间张量那样用完就回收。LSTM 的 hidden state、control flow 子图之间的共享数据都需要这种“耐住性子”的内存。如果误把它当成普通中间张量来复用模型跑两轮状态就被覆盖推理结果直接崩。3.2 PlanAllocations 的核心循环优先吃“回头草”TFLite 的默认内存规划器是ArenaPlanner它在Interpreter::AllocateTensors()时触发规划。整个规划过程大致分四步基于图的拓扑序确定所有节点的执行编号遍历张量结合节点的输入输出关系算出每个张量的出生点和死亡点按节点顺序扫描维护一个“活跃列表”和一个“空闲块池”给每个输出张量分配内存时先扫空闲块池有足够大的块就用它没有就追加新块到 Arena 末尾。这个“优先吃回头草”的策略对运行性能的意义非常大。推理引擎一旦规划完成整个稳定运行阶段几乎不会再触发malloc/free系统调用所有内存都以偏移量落到 Arena 这块连续内存里。这也是为什么同一套模型使用 ArenaPlanner 比逐层手动分配要稳定很多不仅峰值低分配次数也少。有个容易被忽略的点TFLite 在做分配时对相同节点的输入输出会特别处理。因为节点执行时输入输出同时存在所以不会把它们的区间判成可共享。很多自己写 mini 推理引擎的朋友第一次实现生命周期共享时容易在这个地方踩坑区间开闭没算对导致数据被覆盖。3.3 对齐和持久区规划器没说的那些隐藏约束内存规划不只是把大小算对那么简单对齐是另一层约束。CPU 上做 SIMD 运算指针至少要 16 字节对齐某些 ARM 平台和 NPU要求 64 字节甚至更高。如果规划器只按字节大小做连续分配一个 tensor 的起始地址不对齐轻则性能变差重则指令直接异常。所以 ArenaPlanner 在分配时需要显式考虑对齐。常见做法是把每个块的实际大小向上取整到对齐倍数让下一个块天然对齐或者维护 arena 内的 offset按对齐要求做边界修正。代价是多占几个 Padding 字节但换来的是所有数据类型都能安全进行向量化访存。持久区又是另一个维度。状态张量、变量张量不能在推理间隙被回收它们必须待在 Arena 的持久区里。持久区的特点是只增不减除非整个模型被卸载否则不会释放。这部分内存不会参与普通中间张量的复用规划器对待它们的态度更像是在 Arena 里划出一块“固定套间”其他临时张量绕着走。3.4 动态 Shape 为什么让规划器为难ArenaPlanner 能精确分配的前提是所有张量形状在规划时已经确定。可现实中的模型经常有动态形状检测后处理的输出数量随图像内容变化NLP 模型的输入长度不固定语音模型每一帧的上下文也可能变。遇到这种情况TFLite 的规划器只能在规划时跳过这些动态张量把它们放进独立的 dynamic 分配通道。动态张量的大小只有执行到那一刻才知道所以每次都需要临时分配这可能带来额外的延迟和碎片。更麻烦的是一旦动态张量出现在模型中间位置它前后算子的内存复用关系也会被打破Arena 的有效利用率会下降。实际工程里我通常建议能固定 batch 就固定能固定 anchor 生成数量就固定给动态输出预留一个大号静态缓冲区而不是放任它触发动态分配。规划器能帮你省下大部分内存但动态形状还是得依赖使用者的约束。4. 实测与踩坑验证内存规划效果以及最容易翻车的三个场景4.1 看内存规划效果除了 RSS 还能看什么很多人验证内存优化效果只看进程 RSS。但 RSS 是整进程的统计包括代码段、堆、mmap 区域、线程栈干扰因素太多很难直接看出规划器到底省了多少内存。更直接的方式是查看内存规划器本身的输出。TFLite 的规划器接口里暴露了类似arena_size()的方法可以在AllocateTensors()调用之后直接拿到 Arena 当前的总字节数。比如说// 伪代码示意不同版本接口名略有差异 size_t bytes interpreter-arena_size();拿到这个数字再和模型的中间张量总和做一个对比就能算出复用率。你也可以在调试器里打断点看每个 tensor 的 data 指针偏移确认两个存活区间不重叠的张量确实拿到了同一块 arena 地址。这是最直观的验证方式也是理解规划器行为最有效的手段。4.2 翻车场景一频繁 ResizeInputTensor 把规划打回原形TFLite 允许推理时修改输入 shape但每次ResizeInputTensor后都要重新调用AllocateTensors()。这一次调用会把整个规划重新跑一遍。如果代码里每帧都在 resize哪怕只是把同一个尺寸再设置一遍也有可能导致规划失效、Arena 重建所有稳定态都被打破。我见过一个项目视频流输入的分辨率偶尔会切到 1080p、720p 来回横跳结果峰值内存不仅没有降反而因为反复重建 Arena 出现明显卡顿。后来把 resize 统一收敛到初始化阶段推理循环里完全不碰 shape内存和延迟都稳定了。如果你确实需要动态分辨率也尽量把分辨率档次固定成几个档位每个档位预建好对应的 interpreter 或规划而不是让规划器反复算。4.3 翻车场景二多 Subgraph 和持久缓冲把共享内存搞串TFLite 支持 control flow一个主图可以有多个子图。多个 Subgraph 之间往往存在共享的 tensor比如 while 循环里的 carry 变量。如果规划器按每个 Subgraph 独立分配这些跨子图的张量就会在多个 Subgraph 的 Arena 里各存一份白白增加内存。更危险的是持久缓冲。普通中间张量在每个子图结束时可以被释放持久张量却要在整个调用链中保持有效。如果我把一个普通中间张量误存成持久张量内存峰值会被抬高反过来如果把状态张量当普通中间张量处理两轮推理的结果就会互相污染。碰到多 Subgraph 模型我的经验是先理一遍 tensor 的消费链把跨子图的张量单独列出来确认它们的 allocation_type 是持久类型。如果是自己写推理引擎最好的做法是让所有子图共享同一个 Arena并且用“持久区”统一管理跨子图状态而不是各管各的。4.4 翻车场景三Delegate 接管之后一块内存被算了两次TFLite 支持把算子交给 GPU delegate、NPU delegate 执行。Delegate 通常有自己独立的内存空间比如 GPU 的显存或 NPU 的工作 buffer。于是会出现一种尴尬局面CPU 侧的内存规划器已经给某个中间张量分配好了 Arena 块但为了把这个张量喂给 GPU又需要复制一份到 GPU 缓冲同一份数据在两种硬件上存在两份。这不算规划器的 bug但却是实际内存优化的盲区。通常的调法是把 delegate 的输入输出尽量对齐到模型的输入输出 tensor避免多余的“中间人”张量或者使用支持外部缓冲的 delegate让 TFLite 的 tensor 直接指向硬件可访问的地址省掉一次拷贝。测量内存时也要把 GPU/NPU 工作区算进来否则你看到的 CPU 侧内存降得很好看整机峰值却没有任何改善。5. 自己写推理引擎时可以照搬的四步内存规划方案5.1 先把执行顺序排成拓扑序自己实现推理引擎时第一步永远是拿到算子的拓扑执行序。你可以用 Kahn 算法做 BFS也可以用 DFS 做后序输出。得到一条无环的执行序列之后所有算子和张量都能对应到一个整数编号生命周期计算才有基础。这里要特别强调执行顺序必须是稳定的不能在运行时随意改变。一旦算子被并行调度或乱序执行基于生命周期做的缓冲区复用就会失效。如果你的推理引擎有并行调度器内存规划必须和调度策略绑定设计而不是先规划完再并行。5.2 给每个 Tensor 标上出生点和死亡点执行序排好之后遍历所有节点记录每个张量的出生节点和最后消费节点。伪代码可以这样写# exec_order 是按拓扑序编号的节点列表 first_write {} last_read {} for idx, node in enumerate(exec_order): for t in node.outputs: first_write[t] idx for t in node.inputs: last_read[t] idx这里的区间最好用半开区间出生点在节点执行完成后死亡点在最后一个消费节点执行完成后。这样相邻两层的输出张量就能在同一个变更点完成“旧块释放、新块占用”链式网络里可以做到 ping-pong 式的两块激活交替复用峰值内存大幅下降。5.3 扫描式分配先用 free 列表里的旧块有了生命周期接下来就是最核心的分配循环。维护两块结构活跃块列表和空闲块列表。每走到一个节点先处理这个节点输入张量的引用计数计数归零就释放对应块放回空闲列表再为输出张量分配块分配时遍历空闲列表找一个足够大的旧块优先复用。伪代码可以这样表达class Arena { def alloc(size): # 优先复用空闲块 for block in self.free_list: if block.size size: self.free_list.remove(block) self.active_list.append(block) return block.offset # 没有可复用的扩展 arena offset len(self.buffer) self.buffer.extend([0] * size) block Block(offset, size) self.active_list.append(block) return offset def free(offset): block find(self.active_list, offset) self.active_list.remove(block) self.free_list.append(block) }这个 FIRST-FIT 扫描方式简单直接大多数场景下已经够用。不用在规划阶段追求极端最优解图着色再漂亮最后也要落到“空闲块够不够大”“对齐够不够”这些工程约束上。5.4 别忘了对齐、动态 tensor 和持久区这三个例外写完上面的循环你的规划器已经能跑。但要达到能上生产环境的程度三个例外必须处理。第一个是对齐。分配时至少按最大自然对齐取整比如 16 字节如果平台有 64 字节缓存行或 DMA 需求就按 64 字节对齐。块大小向上取整才会不出现越界和对齐错误。第二个是动态 tensor。规划时必须把它们排除在 Arena 稳定分配之外给它们单独一条动态分配路径。如果你强行把动态 tensor 放进静态 Arena大小判断就会失准要么内存浪费要么运行时越界写坏内存。第三个是持久区。跨推理存活的状态张量绝不能参与普通复用。实现时可以在 Arena 里单独开一段持久区只在模型加载时增长推理循环里不回收。否则 LSTM 这类带状态的模型你跑第二次推理时第一次的 hidden state 已经被覆盖了结果完全不可控。5.5 规划策略取舍不一定非要“最优”面对“要不要上一个更优的图着色算法”这个问题我的答案是看场景。策略内存效率实现复杂度运行时开销适用场景FIRST-FIT 顺序扫描高低极低静态图、单次规划、移动端BEST-FIT 精确匹配较高中略高频繁重新规划的动态 shape区间图全局着色最高高离线阶段高编译期离线规划的引擎对绝大多数移动端推理场景FIRST-FIT 已经能拿到接近最优的结果。而且它在运行时几乎不消耗额外 CPU因为规划只在AllocateTensors阶段跑一次。真正应该花时间优化的反而是如何减少重新规划的触发、如何和大块内存池配合减少系统调用。规划策略再精如果每次都重新跑那也是得不偿失。6. 压内存不只靠规划器融合、精度和零拷贝的组合拳6.1 算子融合从源头消灭中间张量规划器是在“中间张量已经存在”的前提下做复用但如果中间张量根本没有那就从源头上省了内存。这就是算子融合的意义。最常见的融合是 ConvReLU、ConvBNReLU。BN 在推理时可以折叠进卷积参数ReLU 作为激活函数可以直接在一个算子内部完成不需要单独产生一个 activation tensor。TFLite 转换器会做很多这类折叠模型转换之后算子数减少内存峰值也自然下降。我见过一些人把融合当成只影响推理速度的优化实际上它对内存影响同样明显。一个 50 层的模型去掉那些只做 Add、ReLU 的独立小算子中间张量数量可能减少三分之一规划器的复用压力小很多。6.2 精度减半arena 直接缩水一半内存大小和精度是线性关系。float32 换成 float16中间张量直接减半换成 int8通常能降到 float32 的四分之一。对量化模型来说卷积输出如果继续保持 int8中间张量的体积会比 float 版小一个量级。这里要注意的是有些算子内部为了精度必须用 float 计算比如某些归一化或更敏感的后处理。模型转 int8 之后这些算子会被插上 Dequantize/Quantize 节点反而多出额外的中间张量。所以量化之后不要只看模型文件变小了也要检查中间张量的变化必要时对个别算子做混合精度处理。6.3 零拷贝让权重和输入根本不进 arena内存规划器管的是中间张量但权重和输入输出这两个环节同样有优化空间。权重能 mmap 就直接 mmap不进 arena 也不进堆输入输出最好让调用方直接提供 buffer解释器把数据写到对应 tensor 的指针里省掉一次外部 vector 到内部 tensor 的 memcpy。如果你自己实现推理引擎可以在 API 层设计一个SetInputTensorData(ptr)的接口而不是要求调用方先准备好一块内存、再拷贝进来。这样既能控制内存又能减少拷贝延迟。中间张量的零拷贝问题相对复杂但如果你的算子支持 in-place 操作比如某些 activation 可以直接在原张量上改写规划器也可以把这个张量标成可复用块省下一整块内存。6.4 一个可参考的实测模型对比为了直观一点给一组来自我实际测试的量级参考。模型是一个 50 层左右的检测模型int8 量化输入 320×320项目数值模型文件大小22 MB所有中间张量大小总和约 180 MB朴素保留全部中间张量的峰值约 180 MBArenaPlanner 规划后的峰值约 26 MB复用后节省约 85%另一类更轻量的分类模型比如 MobileNetV2 1.0模型文件 14MB中间张量总和约 45MB规划后峰值不到 15MB。不同网络结构差异很大但规律一致权重占比低、激活占比高的模型内存规划器带来的收益最明显。6.5 内存还是紧从这三个方向接着抠如果规划器已经跑满内存还是吃紧我一般从三个方向继续抠。第一降分辨率或固定 batch1这直接影响 activation map 的宽度和高度收益立刻可见第二用子图拆分把大模型切成几段每段单独分配内存让峰值窗口缩小第三检查后处理和输入预处理很多内存峰值不是推理引擎造成的而是图像解码、NMS 结果堆叠这些外围逻辑悄悄吃掉的。最后分享一点个人体会把内存规划做扎实之后最舒服的不是峰值降了多少而是内存行为变得可预测。单次推理的 malloc 次数从几百次降到个位数RSS 曲线在推理循环里基本是一条直线AllocateTensors 只出现在初始化阶段。这种“一切尽在掌控”的稳定感才是推理引擎真正可以上生产线的底气。内存规划器这个管家省的不只是字节还有你在排查线上 OOM 时快要耗尽的耐心。