内存又超了——这句话我几乎每次在端侧模型压测会上都能听到。模型文件明明只有几百 KB一加载实测峰值却冲到几百 MB或者改一下输入分辨率App 还没推理完就被系统按内存上限给杀掉了。很多人第一反应是去查模型量化没量化、算子有没有冗余其实还有一个底层非常容易被忽略的环节推理引擎里的内存规划器Memory Planner。在 TFLite 里它专门负责调度中间张量的内存复用相当于整个推理管线的“内存管家”。这篇文章我打算把这个管家的工作机制完整拆开讲一遍它到底在管些什么、是怎么分配和复用的、省内存的实际效果能到多少以及我在真实部署项目里踩过的坑和调试方法。对做端侧部署、跑 TFLite 在 Android/iOS/嵌入式平台上、或者单纯想把模型内存峰值压下去的人这个方向值得从头到尾过一遍。1. 没有“内存管家”的推理引擎会怎样朴素分配的内存账本1.1 朴素分配方式每个中间张量都有一块私有缓冲区推理引擎执行一张模型图时不只有模型权重占内存。每个算子从输入张量算到输出张量输出张量通常叫中间激活要占用一块缓冲区如果网络很深层层堆叠下来中间张量的数量会非常可观。最朴素的实现方式是每遇到一个中间张量就malloc一块大小刚好等于tensor-bytes的内存算子执行完也不急着释放等到模型跑完再一次清理。这样做的结果就是内存峰值约等于所有中间张量大小之和。以我经常跑的 MobileNetV2 1.0 为例224x224 的 RGB 输入浮点模型所有中间激活累加起来在没有规划的情况下需要几十 MB 级别的一块“叠加总量”。对服务端不算什么但对手机尤其是低内存机型来说这基本属于不可接受的起步开销。更要命的是如果模型里多个分支并行存在每个分支的中间张量可能同时活着内存前缀和会进一步膨胀OOM 几乎是必然的。1.2 张量生命周期一张由执行顺序决定的“入住时间表”要打破“各占一块”的笨办法先要回答一个关键问题每个中间张量到底在哪个时间段需要内存这个时间段不是看模型拓扑顺序而是看实际执行顺序。TFLite 在加载模型后会生成一个执行计划把图中的算子排成一个线性数组这个顺序通常是拓扑排序的结果。某个算子在执行计划里的下标决定了它所有输入张量的“最后使用节点”和输出张量的“产生节点”。于是每个张量都能算出自己的生命周期区间[producer_node, last_consumer_node]张量被某个算子作为输出产生的那一个节点是它的“出生点”。张量作为输入被最后一个算子消费掉的那个节点是它的“死亡点”。只要在这个区间内这个张量就必须占着内存谁都动不了它区间一旦结束它所占的内存就可以交给别人了。这里有一个容易被忽略的细节一个张量即使看起来在整个模型推理过程中“一直都在”它是否真的全程占用内存取决于它被最后一个消费的节点排在哪里。如果最后一层才用到它那它的生命周期会横跨大半个执行计划这种长生命周期张量就是内存规划里最难处理的“钉子户”后面我会专门展开说。1.3 内存复用的核心思想错峰入住的宿舍分配有了生命周期表内存规划的目标就变成把所有张量按生命周期塞进一块总内存里让生命周期不重叠的张量共用同一块地址空间。我习惯把它想成宿舍分配——每个人有明确的入住和退房时间只要错峰就可以让不同的人住进同一个床位宿舍房间总数就能远小于入住总人数。举一个能说明问题的微型例子。假设执行计划是连续 4 个节点张量情况如下张量产生节点最后使用节点大小T11216 KBT2234 KBT3348 KBT4452 KB朴素方案需要 16482 30KB。但看生命周期T1 到节点 2 就结束T3 从节点 3 才开始它俩完全不重叠T3 可以直接睡进 T1 退出来的那 16KBT2 到节点 3 结束T4 恰好从节点 4 才出现虽然 T4 只要 2KB但 4KB 的床位也能复用。实际峰值出现在节点 2T116KB T24KB同时活着总共 20KB比朴素方案省了三分之一。真实模型里张量数量多、生命周期交错复杂人工看表不现实所以这个任务就交给内存规划器去做而它在 TFLite 里的实现核心就是一个基于贪心策略的分配流程。这就是“内存管家”存在的原因。2. TFLite 内存规划器的实现拆解从 ArenaPlanner 到贪心分配2.1 Arena 分配器一次性申请、按需切分的思路很多人会误以为 TFLite 推理时每个算子都在动态malloc。实际上推理期间基本不会碰堆分配器。调用AllocateTensors()时规划器先算出每个张量在内存里的偏移量然后一次性申请一个大缓冲区可能几 MB 到几十 MB之后所有中间张量都在这块区域里切分、复用。这就是 Arena内存竞技场分配器SimpleMemoryArena干的就是这件事。Arena 带来的好处很直接减少了大量小碎片的malloc/free内存布局稳定甚至对 Cache 命中率都有帮助。推理引擎在热循环里如果频繁分配释放不仅慢还容易出现堆碎片导致 OOM。Arena 把问题从“堆管理”变成了“规划偏移量”简单、可控、且可复现。2.2 AllocationInfo 与执行计划的关系在ArenaPlanner内部每个可变张量都会对应一条AllocationInfo记录里面有三个最核心的整数字段产生该张量的节点下标、第一个消费该张量的节点下标、最后一个消费该张量的节点下标。这三个值一起刻画了张量的生命周期窗口。接下来规划器会做一轮“身份筛查”模型输入/输出张量指针不能挪动必须单独处理。变量张量is_variable true比如 RNN 状态、持久缓存值要在多次Invoke之间保留不可复用。其余普通中间张量才进入可复用池。这条筛查逻辑很重要。我做调试时发现很多人把“内存复用”理解为所有张量都可以随便重叠结果自定义算子把该持久的数据写进可复用区下一次推理时数据就被覆盖了整个模型输出错乱。这个问题不是规划器导致的而是对生命周期定义不清楚。2.3 贪心策略的分配流程与偏移计算拿到清晰的 AllocationInfo 后规划器开始分配大致流程如下按执行计划从左到右扫描各个节点。每到一个节点检查所有已“退房”的空闲块即最后使用节点已经小于当前节点的张量所占空间。为当前节点需要的新张量找一块足够大的空闲块如果有就复用没有就向 arena 尾部追加。新产生的输出张量登记为“在住”生命周期窗口继续往前走。全部节点扫描完毕每个张量都得到一个相对于 arena 基址的偏移量。这里的具体选择策略比如优先找最大块还是最合适的块在不同版本里有差异有的版本用 first-fit 思想有的版本倾向于大块优先。它本质是一个贪心近似并不保证全局最优但胜在速度快、实现简单。几百个张量的扫描在加载阶段几毫秒内就能完成完全不会影响推理时延。2.4 64 字节对齐为了让 SIMD 指令“跑得动”内存规划不是简单地把大小相加后塞洞它还必须满足对齐要求。TFLite 的大量 CPU 算子会使用 NEON / SSE / AVX 这类 SIMD 指令做向量化加载这些指令往往要求地址按 16 字节或更大粒度对齐。如果规划器把一个小张量随便塞进上一个张量留下的空洞起始地址不对齐轻则 kernel 崩溃重则只能退回到慢速标量路径内存省了但性能崩了。所以 arena 在分配时会把每个张量的起始地址和占用大小都向上取整到对齐边界。TFLite 里常见的默认对齐值是 64 字节具体以自己手上版本的源码为准。这也是为什么规划后的 arena 总大小往往会比理论计算值略大一点这点“冗余”是为了让加速指令跑得稳值得接受。2.5 持久张量内存规划管不了的那部分并非所有内存都进可复用的 arena。TFLite 里大致分三类权重张量来自.tflite文件通常走 mmap 只读映射根本不参与规划这就是为什么模型文件大小不直接等于运行内存。模型输入/输出张量用户代码随时会引用和修改规划器必须保证它们的指针稳定不做复用。变量张量跨Invoke保持状态也不能被挪动或覆盖。规划器会把它们标记为 persistent独立管理。实际部署时如果你发现内存比预期高出一截可以先检查是不是模型里变量张量或者 delegate 边界张量太多把本该复用的大块空间给“钉”住了——这种情况在 RNN 和 Transformer 类模型里尤其常见。3. 实测数据说话不同模型上的内存规划效果3.1 MobileNet 这种轻量模型中间张量竟然这么多MobileNet 系列已是出了名的轻量参数只有几 MB但中间激活一点都不少。我在压测环境里用浮点 MobileNetV2 1.0 跑 224x224 输入时中间张量数量大概在 200 个上下。把所有 tensor 的bytes全部加起来的朴素合计能达到几十 MB经过规划后的 arena 实际峰值通常会压到十几 MB 的级别节省幅度在 4~5 倍左右。这还只是 224x224分辨率提到 512x512 时中间特征图的面积变成 4 倍大张量的生命周期一旦重叠峰值会肉眼可见地涨规划器能腾挪的空间也随之变化。所以别因为模型是“轻量级”就忽视内存规划的作用轻量参数的模型照样可以因为激活张量拖垮内存。3.2 BERT 类模型激活张量与残差连接的“钉子户”NLP 模型的情况不太一样。以 BERT 类模型为例序列长度直接影响 attention 矩阵大小。假设 batch1、序列长度 128、隐藏维度 512一层 attention 分数矩阵大约是 128x128 乘以头数个浮点数单层不算大但多层累积起来再加上 FFN 的中间层和层间残差连接整张生命周期表会被拉得很长。好消息是BERT 每层几乎是顺序执行的层 i 的大部分中间张量在层 i1 开始前就到期了所以层间复用非常充分。这也是为什么很多端侧 BERT 实际激活内存远小于“把所有层全叠加”的理论值。坏消息是残差连接相当于在每个 Transformer 层里埋了一条贯穿整个 block 的生命周期如果执行计划和图优化没把残差处理好这些长生命周期张量就会把峰值钉在高位。工作里我确实遇到过同一结构的模型换了 TFLite 版本之后内存直接差出近三成的情况排查下来问题就出在残差张量生命周期处理策略的差异上。3.3 检测模型多分支并行让规划器更难发挥反例是 SSD、特征金字塔这类多尺度检测模型。多分支结构有一个特点多个分支的张量需要在汇合点之前一直活着分支越宽、汇合点越靠后同一时刻活着的张量就越多规划器能做的“错峰”就越有限。这时候峰值更多由网络结构决定而不是规划策略。对这类模型单纯指望规划器还不如回头改结构。例如把多个汇合点拆开先让一个分支跑完释放内存再做另一个分支或者把多尺度特征提前降采样减少并行分支的宽大张量。效果往往立竿见影。这也是我常和算法同学讲的一句话内存优化不应该只是部署阶段的补救而是结构设计阶段就要考虑的事情。3.4 一次规划、多次使用的静态特性对推理时延的影响内存规划是一次性动作。AllocateTensors()期间完成所有计算和内存分配之后每次Invoke()只是照着既定偏移执行算子不会再做任何malloc、计算偏移量的操作。所以规划器对单次推理时延的影响基本可以忽略对首次加载的影响也通常在毫秒级几百个张量的贪心扫描其实非常快。真正要注意的是如果推理过程中动态改变输入尺寸系统会重新规划才可能出现可感知的卡顿或内存抖动这个问题放到下一节展开。4. 我在部署时踩过的坑动态形状、ScratchBuffer 与调试4.1 动态输入尺寸每次 Resize 都会推倒重来TFLite 里动态改变输入尺寸会调用ResizeInputTensor接着下一次Invoke时因为部分中间张量的大小也跟着变了原来的内存计划失效规划器要重新计算。功能上没问题但有两个后果一是峰值抖动。重规划时旧 arena 可能还没释放新 arena 又建出来一瞬间内存可能翻倍。如果业务侧高频切换输入分辨率比如相机预览从 640x480 切到 1280x720非常容易触发 OOM。二是碎片化风险。反复重规划后arena 里会留下各种尺寸的洞贪心策略不一定能填好实际使用量可能一次比一次大。我踩过的具体场景是做相机预览的人脸检测预览尺寸跟随手机旋转变化模型被频繁 Resize最后在低内存手机上稳定复现 OOM。后来的处理方法是固定允许的输入尺寸集合启动时直接按最大尺寸规划一次推理时把输入图片缩放裁切到模型实际需要的大小保持模型输入形状不变。这样内存计划只建一次峰值完全可控。4.2 算子 Scratch Buffer内存账本里最容易被忽略的一行TFLite 的算子可以通过RequestScratchBufferInArena向 context 申请临时缓冲区。这个机制常见于 GEMM 类算子Conv、FullyConnected的优化 kernel它们一次性把中间计算用的“草稿纸”放进 arena避免每次调用临时malloc。这个设计本身没问题但自定义算子或者新版本 kernel 的 scratch 需求经常超出预期。我遇到过一件真实的事升级 TFLite 版本后某个卷积算子的 scratch 大小翻了接近一倍导致 arena 内存直接上涨但模型张量本身大小完全没变。如果只盯着tensor-bytes查永远查不到原因。所以排查内存问题时要把每个算子的 scratch buffer 也翻出来看。如果发现某个算子占的 scratch 过大可以考虑换一个内存友好的 kernel 实现或者把输入切小、用算子融合绕开那块临时区。4.3 怎么确认内存真的被复用一条可行的调试路径判断规划器到底有没有干活最直接的方法是看张量的内存地址。规划生效后生命周期不重叠的张量会拿到同一个data.raw指针同一个偏移量。我在调试时经常写这样一段小工具// 打印所有张量的指针与大小用于判断内存复用情况 for (int i 0; i interpreter-tensors_size(); i) { TfLiteTensor* t interpreter-tensor(i); if (t-data.raw nullptr) continue; printf(tensor[%d] ptr%p bytes%zu\n, i, static_castvoid*(t-data.raw), t-bytes); }如果两个生命周期明显不重叠的张量打印出来的指针一样说明规划器成功复用了如果每个张量指针都不一样要么是你的模型几乎全程高并发要么是规划信息被破坏了比如自定义算子把生命周期信息弄错。另外也可以通过指针地址判断对齐情况打印出来的地址如果没有按 16 或 64 字节对齐那基本可以断定某个自定义路径绕过了 arena 的对齐逻辑SIMD kernel 很容易在这里翻车。4.4 多线程与 GPU/NPU 场景下的内存叠加效应多线程推理不会改变内存规划的逻辑因为规划基于执行计划而不是线程数量。但线程池和算子内部并发临时 buffer 会额外占用内存。如果开 4 线程跑同一个模型各个算子内部可能各自准备一份 scratch这部分不在 arena 规划范围内。GPU/NPU 委托delegate则更明显CPU 侧的 arena 只负责 CPU 子图和输入输出桥接GPU 侧自己在运行时申请缓冲区。委托切分图的边界张量需要固定指针以便跨后端拷贝这些张量会被标记为 persistent进而抬高 CPU arena 的使用。我见过一种“诡异”情况把一个卷积层委托给 GPU 后CPU 侧内存反而涨了一截原因就是边界 persistent 张量变多了。所以做委托时要合并考虑两侧内存不能只盯着帧率提升那一项。5. 几个真正提升内存效率的部署建议从规划器之外下功夫5.1 量化带来的内存连锁反应量化是把 float32 降到 int8 或 float16张量 bytes 直接缩到原来的 1/4 或 1/2。arena 里绝大多数是激活张量它们全部按比例缩小因此峰值内存几乎同比例下降。这对内存的改善比任何规划器调参都直接。但要注意混合量化场景某些算子仍然需要 float32 输入输出会产生 dequant/requant 的临时 float 张量。这些 float 张量生命周期可能很短规划器有机会复用但如果它们分布在每个算子附近会让 arena 里同时存在“int8 密集区”和“float 稀疏区”实际节省不一定能打到理论上的 1/4。部署时可以先看量化前后各张量bytes的分布再决定要不要做算子级 float 融合。5.2 图优化与算子融合先给规划器“减负”规划器只能在给定生命周期表里找复用空间它改不了生命周期本身。如果某些长生命周期张量只是结构性产物——比如没有实际作用的 identity、可以吸收进前一个算子的 BatchNorm、Add那么通过图优化把它们去掉生命周期表会明显缩短规划器能腾挪的空间自然变大。我常做的检查是模型导出 tflite 之前先看有没有可以融合的算子组合导出之后再用工具对比融合前后的生命周期分布。很多时候内存下降 20%~30% 不是靠调什么参数而是靠让没必要的张量直接消失。这个思路在检测模型里尤其好用因为多分支结构天然容易产生大量过渡张量。5.3 委托分配之后CPU 侧内存反而变大的现象这是上一节“边界 persistent 张量”问题的延伸思考。委托给 GPU/NPU 后模型被切成若干子图每个子图的输入输出张量要在 CPU 侧保持驻留供跨后端拷贝。原本一个横跨整个网络的中间张量只要位于子图边界就会被升级为 persistent不再参与复用。子图切得越碎这类固定张量就越多CPU arena 反而可能上涨。所以实际落地时我会建议先看整体子图划分方案尽量把神经网络中段整体委托给同一后端减少跨后端的边界如果 CPU 侧内存本来不是关键路径甚至可以主动接受部分上涨换取 GPU 整体推理时间的下降。内存规划和硬件调度是同一盘棋不能只盯着一边看。根据我个人的实际部署体会内存规划器的价值往往要在量化、图融合、委托切分全部做完之后才真正体现出来。它是托住模型运行高峰的那只手但手能有多大空间很大程度取决于你上游把模型加工得多干净。真到了低内存设备上做最后冲刺时与其反复怀疑规划器不如回到那张生命周期表把“钉子户”张量一个个找出来干掉。
TFLite内存规划器深度拆解：从张量生命周期到端侧内存峰值优化
内存又超了——这句话我几乎每次在端侧模型压测会上都能听到。模型文件明明只有几百 KB一加载实测峰值却冲到几百 MB或者改一下输入分辨率App 还没推理完就被系统按内存上限给杀掉了。很多人第一反应是去查模型量化没量化、算子有没有冗余其实还有一个底层非常容易被忽略的环节推理引擎里的内存规划器Memory Planner。在 TFLite 里它专门负责调度中间张量的内存复用相当于整个推理管线的“内存管家”。这篇文章我打算把这个管家的工作机制完整拆开讲一遍它到底在管些什么、是怎么分配和复用的、省内存的实际效果能到多少以及我在真实部署项目里踩过的坑和调试方法。对做端侧部署、跑 TFLite 在 Android/iOS/嵌入式平台上、或者单纯想把模型内存峰值压下去的人这个方向值得从头到尾过一遍。1. 没有“内存管家”的推理引擎会怎样朴素分配的内存账本1.1 朴素分配方式每个中间张量都有一块私有缓冲区推理引擎执行一张模型图时不只有模型权重占内存。每个算子从输入张量算到输出张量输出张量通常叫中间激活要占用一块缓冲区如果网络很深层层堆叠下来中间张量的数量会非常可观。最朴素的实现方式是每遇到一个中间张量就malloc一块大小刚好等于tensor-bytes的内存算子执行完也不急着释放等到模型跑完再一次清理。这样做的结果就是内存峰值约等于所有中间张量大小之和。以我经常跑的 MobileNetV2 1.0 为例224x224 的 RGB 输入浮点模型所有中间激活累加起来在没有规划的情况下需要几十 MB 级别的一块“叠加总量”。对服务端不算什么但对手机尤其是低内存机型来说这基本属于不可接受的起步开销。更要命的是如果模型里多个分支并行存在每个分支的中间张量可能同时活着内存前缀和会进一步膨胀OOM 几乎是必然的。1.2 张量生命周期一张由执行顺序决定的“入住时间表”要打破“各占一块”的笨办法先要回答一个关键问题每个中间张量到底在哪个时间段需要内存这个时间段不是看模型拓扑顺序而是看实际执行顺序。TFLite 在加载模型后会生成一个执行计划把图中的算子排成一个线性数组这个顺序通常是拓扑排序的结果。某个算子在执行计划里的下标决定了它所有输入张量的“最后使用节点”和输出张量的“产生节点”。于是每个张量都能算出自己的生命周期区间[producer_node, last_consumer_node]张量被某个算子作为输出产生的那一个节点是它的“出生点”。张量作为输入被最后一个算子消费掉的那个节点是它的“死亡点”。只要在这个区间内这个张量就必须占着内存谁都动不了它区间一旦结束它所占的内存就可以交给别人了。这里有一个容易被忽略的细节一个张量即使看起来在整个模型推理过程中“一直都在”它是否真的全程占用内存取决于它被最后一个消费的节点排在哪里。如果最后一层才用到它那它的生命周期会横跨大半个执行计划这种长生命周期张量就是内存规划里最难处理的“钉子户”后面我会专门展开说。1.3 内存复用的核心思想错峰入住的宿舍分配有了生命周期表内存规划的目标就变成把所有张量按生命周期塞进一块总内存里让生命周期不重叠的张量共用同一块地址空间。我习惯把它想成宿舍分配——每个人有明确的入住和退房时间只要错峰就可以让不同的人住进同一个床位宿舍房间总数就能远小于入住总人数。举一个能说明问题的微型例子。假设执行计划是连续 4 个节点张量情况如下张量产生节点最后使用节点大小T11216 KBT2234 KBT3348 KBT4452 KB朴素方案需要 16482 30KB。但看生命周期T1 到节点 2 就结束T3 从节点 3 才开始它俩完全不重叠T3 可以直接睡进 T1 退出来的那 16KBT2 到节点 3 结束T4 恰好从节点 4 才出现虽然 T4 只要 2KB但 4KB 的床位也能复用。实际峰值出现在节点 2T116KB T24KB同时活着总共 20KB比朴素方案省了三分之一。真实模型里张量数量多、生命周期交错复杂人工看表不现实所以这个任务就交给内存规划器去做而它在 TFLite 里的实现核心就是一个基于贪心策略的分配流程。这就是“内存管家”存在的原因。2. TFLite 内存规划器的实现拆解从 ArenaPlanner 到贪心分配2.1 Arena 分配器一次性申请、按需切分的思路很多人会误以为 TFLite 推理时每个算子都在动态malloc。实际上推理期间基本不会碰堆分配器。调用AllocateTensors()时规划器先算出每个张量在内存里的偏移量然后一次性申请一个大缓冲区可能几 MB 到几十 MB之后所有中间张量都在这块区域里切分、复用。这就是 Arena内存竞技场分配器SimpleMemoryArena干的就是这件事。Arena 带来的好处很直接减少了大量小碎片的malloc/free内存布局稳定甚至对 Cache 命中率都有帮助。推理引擎在热循环里如果频繁分配释放不仅慢还容易出现堆碎片导致 OOM。Arena 把问题从“堆管理”变成了“规划偏移量”简单、可控、且可复现。2.2 AllocationInfo 与执行计划的关系在ArenaPlanner内部每个可变张量都会对应一条AllocationInfo记录里面有三个最核心的整数字段产生该张量的节点下标、第一个消费该张量的节点下标、最后一个消费该张量的节点下标。这三个值一起刻画了张量的生命周期窗口。接下来规划器会做一轮“身份筛查”模型输入/输出张量指针不能挪动必须单独处理。变量张量is_variable true比如 RNN 状态、持久缓存值要在多次Invoke之间保留不可复用。其余普通中间张量才进入可复用池。这条筛查逻辑很重要。我做调试时发现很多人把“内存复用”理解为所有张量都可以随便重叠结果自定义算子把该持久的数据写进可复用区下一次推理时数据就被覆盖了整个模型输出错乱。这个问题不是规划器导致的而是对生命周期定义不清楚。2.3 贪心策略的分配流程与偏移计算拿到清晰的 AllocationInfo 后规划器开始分配大致流程如下按执行计划从左到右扫描各个节点。每到一个节点检查所有已“退房”的空闲块即最后使用节点已经小于当前节点的张量所占空间。为当前节点需要的新张量找一块足够大的空闲块如果有就复用没有就向 arena 尾部追加。新产生的输出张量登记为“在住”生命周期窗口继续往前走。全部节点扫描完毕每个张量都得到一个相对于 arena 基址的偏移量。这里的具体选择策略比如优先找最大块还是最合适的块在不同版本里有差异有的版本用 first-fit 思想有的版本倾向于大块优先。它本质是一个贪心近似并不保证全局最优但胜在速度快、实现简单。几百个张量的扫描在加载阶段几毫秒内就能完成完全不会影响推理时延。2.4 64 字节对齐为了让 SIMD 指令“跑得动”内存规划不是简单地把大小相加后塞洞它还必须满足对齐要求。TFLite 的大量 CPU 算子会使用 NEON / SSE / AVX 这类 SIMD 指令做向量化加载这些指令往往要求地址按 16 字节或更大粒度对齐。如果规划器把一个小张量随便塞进上一个张量留下的空洞起始地址不对齐轻则 kernel 崩溃重则只能退回到慢速标量路径内存省了但性能崩了。所以 arena 在分配时会把每个张量的起始地址和占用大小都向上取整到对齐边界。TFLite 里常见的默认对齐值是 64 字节具体以自己手上版本的源码为准。这也是为什么规划后的 arena 总大小往往会比理论计算值略大一点这点“冗余”是为了让加速指令跑得稳值得接受。2.5 持久张量内存规划管不了的那部分并非所有内存都进可复用的 arena。TFLite 里大致分三类权重张量来自.tflite文件通常走 mmap 只读映射根本不参与规划这就是为什么模型文件大小不直接等于运行内存。模型输入/输出张量用户代码随时会引用和修改规划器必须保证它们的指针稳定不做复用。变量张量跨Invoke保持状态也不能被挪动或覆盖。规划器会把它们标记为 persistent独立管理。实际部署时如果你发现内存比预期高出一截可以先检查是不是模型里变量张量或者 delegate 边界张量太多把本该复用的大块空间给“钉”住了——这种情况在 RNN 和 Transformer 类模型里尤其常见。3. 实测数据说话不同模型上的内存规划效果3.1 MobileNet 这种轻量模型中间张量竟然这么多MobileNet 系列已是出了名的轻量参数只有几 MB但中间激活一点都不少。我在压测环境里用浮点 MobileNetV2 1.0 跑 224x224 输入时中间张量数量大概在 200 个上下。把所有 tensor 的bytes全部加起来的朴素合计能达到几十 MB经过规划后的 arena 实际峰值通常会压到十几 MB 的级别节省幅度在 4~5 倍左右。这还只是 224x224分辨率提到 512x512 时中间特征图的面积变成 4 倍大张量的生命周期一旦重叠峰值会肉眼可见地涨规划器能腾挪的空间也随之变化。所以别因为模型是“轻量级”就忽视内存规划的作用轻量参数的模型照样可以因为激活张量拖垮内存。3.2 BERT 类模型激活张量与残差连接的“钉子户”NLP 模型的情况不太一样。以 BERT 类模型为例序列长度直接影响 attention 矩阵大小。假设 batch1、序列长度 128、隐藏维度 512一层 attention 分数矩阵大约是 128x128 乘以头数个浮点数单层不算大但多层累积起来再加上 FFN 的中间层和层间残差连接整张生命周期表会被拉得很长。好消息是BERT 每层几乎是顺序执行的层 i 的大部分中间张量在层 i1 开始前就到期了所以层间复用非常充分。这也是为什么很多端侧 BERT 实际激活内存远小于“把所有层全叠加”的理论值。坏消息是残差连接相当于在每个 Transformer 层里埋了一条贯穿整个 block 的生命周期如果执行计划和图优化没把残差处理好这些长生命周期张量就会把峰值钉在高位。工作里我确实遇到过同一结构的模型换了 TFLite 版本之后内存直接差出近三成的情况排查下来问题就出在残差张量生命周期处理策略的差异上。3.3 检测模型多分支并行让规划器更难发挥反例是 SSD、特征金字塔这类多尺度检测模型。多分支结构有一个特点多个分支的张量需要在汇合点之前一直活着分支越宽、汇合点越靠后同一时刻活着的张量就越多规划器能做的“错峰”就越有限。这时候峰值更多由网络结构决定而不是规划策略。对这类模型单纯指望规划器还不如回头改结构。例如把多个汇合点拆开先让一个分支跑完释放内存再做另一个分支或者把多尺度特征提前降采样减少并行分支的宽大张量。效果往往立竿见影。这也是我常和算法同学讲的一句话内存优化不应该只是部署阶段的补救而是结构设计阶段就要考虑的事情。3.4 一次规划、多次使用的静态特性对推理时延的影响内存规划是一次性动作。AllocateTensors()期间完成所有计算和内存分配之后每次Invoke()只是照着既定偏移执行算子不会再做任何malloc、计算偏移量的操作。所以规划器对单次推理时延的影响基本可以忽略对首次加载的影响也通常在毫秒级几百个张量的贪心扫描其实非常快。真正要注意的是如果推理过程中动态改变输入尺寸系统会重新规划才可能出现可感知的卡顿或内存抖动这个问题放到下一节展开。4. 我在部署时踩过的坑动态形状、ScratchBuffer 与调试4.1 动态输入尺寸每次 Resize 都会推倒重来TFLite 里动态改变输入尺寸会调用ResizeInputTensor接着下一次Invoke时因为部分中间张量的大小也跟着变了原来的内存计划失效规划器要重新计算。功能上没问题但有两个后果一是峰值抖动。重规划时旧 arena 可能还没释放新 arena 又建出来一瞬间内存可能翻倍。如果业务侧高频切换输入分辨率比如相机预览从 640x480 切到 1280x720非常容易触发 OOM。二是碎片化风险。反复重规划后arena 里会留下各种尺寸的洞贪心策略不一定能填好实际使用量可能一次比一次大。我踩过的具体场景是做相机预览的人脸检测预览尺寸跟随手机旋转变化模型被频繁 Resize最后在低内存手机上稳定复现 OOM。后来的处理方法是固定允许的输入尺寸集合启动时直接按最大尺寸规划一次推理时把输入图片缩放裁切到模型实际需要的大小保持模型输入形状不变。这样内存计划只建一次峰值完全可控。4.2 算子 Scratch Buffer内存账本里最容易被忽略的一行TFLite 的算子可以通过RequestScratchBufferInArena向 context 申请临时缓冲区。这个机制常见于 GEMM 类算子Conv、FullyConnected的优化 kernel它们一次性把中间计算用的“草稿纸”放进 arena避免每次调用临时malloc。这个设计本身没问题但自定义算子或者新版本 kernel 的 scratch 需求经常超出预期。我遇到过一件真实的事升级 TFLite 版本后某个卷积算子的 scratch 大小翻了接近一倍导致 arena 内存直接上涨但模型张量本身大小完全没变。如果只盯着tensor-bytes查永远查不到原因。所以排查内存问题时要把每个算子的 scratch buffer 也翻出来看。如果发现某个算子占的 scratch 过大可以考虑换一个内存友好的 kernel 实现或者把输入切小、用算子融合绕开那块临时区。4.3 怎么确认内存真的被复用一条可行的调试路径判断规划器到底有没有干活最直接的方法是看张量的内存地址。规划生效后生命周期不重叠的张量会拿到同一个data.raw指针同一个偏移量。我在调试时经常写这样一段小工具// 打印所有张量的指针与大小用于判断内存复用情况 for (int i 0; i interpreter-tensors_size(); i) { TfLiteTensor* t interpreter-tensor(i); if (t-data.raw nullptr) continue; printf(tensor[%d] ptr%p bytes%zu\n, i, static_castvoid*(t-data.raw), t-bytes); }如果两个生命周期明显不重叠的张量打印出来的指针一样说明规划器成功复用了如果每个张量指针都不一样要么是你的模型几乎全程高并发要么是规划信息被破坏了比如自定义算子把生命周期信息弄错。另外也可以通过指针地址判断对齐情况打印出来的地址如果没有按 16 或 64 字节对齐那基本可以断定某个自定义路径绕过了 arena 的对齐逻辑SIMD kernel 很容易在这里翻车。4.4 多线程与 GPU/NPU 场景下的内存叠加效应多线程推理不会改变内存规划的逻辑因为规划基于执行计划而不是线程数量。但线程池和算子内部并发临时 buffer 会额外占用内存。如果开 4 线程跑同一个模型各个算子内部可能各自准备一份 scratch这部分不在 arena 规划范围内。GPU/NPU 委托delegate则更明显CPU 侧的 arena 只负责 CPU 子图和输入输出桥接GPU 侧自己在运行时申请缓冲区。委托切分图的边界张量需要固定指针以便跨后端拷贝这些张量会被标记为 persistent进而抬高 CPU arena 的使用。我见过一种“诡异”情况把一个卷积层委托给 GPU 后CPU 侧内存反而涨了一截原因就是边界 persistent 张量变多了。所以做委托时要合并考虑两侧内存不能只盯着帧率提升那一项。5. 几个真正提升内存效率的部署建议从规划器之外下功夫5.1 量化带来的内存连锁反应量化是把 float32 降到 int8 或 float16张量 bytes 直接缩到原来的 1/4 或 1/2。arena 里绝大多数是激活张量它们全部按比例缩小因此峰值内存几乎同比例下降。这对内存的改善比任何规划器调参都直接。但要注意混合量化场景某些算子仍然需要 float32 输入输出会产生 dequant/requant 的临时 float 张量。这些 float 张量生命周期可能很短规划器有机会复用但如果它们分布在每个算子附近会让 arena 里同时存在“int8 密集区”和“float 稀疏区”实际节省不一定能打到理论上的 1/4。部署时可以先看量化前后各张量bytes的分布再决定要不要做算子级 float 融合。5.2 图优化与算子融合先给规划器“减负”规划器只能在给定生命周期表里找复用空间它改不了生命周期本身。如果某些长生命周期张量只是结构性产物——比如没有实际作用的 identity、可以吸收进前一个算子的 BatchNorm、Add那么通过图优化把它们去掉生命周期表会明显缩短规划器能腾挪的空间自然变大。我常做的检查是模型导出 tflite 之前先看有没有可以融合的算子组合导出之后再用工具对比融合前后的生命周期分布。很多时候内存下降 20%~30% 不是靠调什么参数而是靠让没必要的张量直接消失。这个思路在检测模型里尤其好用因为多分支结构天然容易产生大量过渡张量。5.3 委托分配之后CPU 侧内存反而变大的现象这是上一节“边界 persistent 张量”问题的延伸思考。委托给 GPU/NPU 后模型被切成若干子图每个子图的输入输出张量要在 CPU 侧保持驻留供跨后端拷贝。原本一个横跨整个网络的中间张量只要位于子图边界就会被升级为 persistent不再参与复用。子图切得越碎这类固定张量就越多CPU arena 反而可能上涨。所以实际落地时我会建议先看整体子图划分方案尽量把神经网络中段整体委托给同一后端减少跨后端的边界如果 CPU 侧内存本来不是关键路径甚至可以主动接受部分上涨换取 GPU 整体推理时间的下降。内存规划和硬件调度是同一盘棋不能只盯着一边看。根据我个人的实际部署体会内存规划器的价值往往要在量化、图融合、委托切分全部做完之后才真正体现出来。它是托住模型运行高峰的那只手但手能有多大空间很大程度取决于你上游把模型加工得多干净。真到了低内存设备上做最后冲刺时与其反复怀疑规划器不如回到那张生命周期表把“钉子户”张量一个个找出来干掉。
