TFLite内存规划器解析：张量生命周期与Arena复用机制

做端侧推理的人，早晚会在内存上栽一回。模型文件 20MB，你觉得稳了，结果一跑起来进程直接飙到 200MB 以上，换个小内存的设备就 OOM。多数人第一反应是换更小的模型、降分辨率，其实还有一个比模型体积更隐蔽的“内存黑洞”——推理引擎内部的中间张量（tensor）。TFLite 里负责给这些张量安排内存位置、想尽办法让它们共用空间的组件，就是内存规划器（Memory Planner）。它在推理引擎里扮演的正是“内存管家”的角色：谁先进场、谁先退场、谁的位子可以腾给谁，都由它决定。

这篇文章想把我对这套机制的理解、调试它时用到的工具，以及踩过的坑整理出来。适合正在做移动端/嵌入式模型部署、遇到过内存飙升或 OOM 的开发者参考。我不打算把源码逐行讲一遍，而是从“内存账本怎么算”“规划器怎么安排”“我实际怎么观测和优化”三个角度展开，最后附上能直接复现的 Python 脚本和排查清单。

1. 为什么推理引擎需要一个“内存管家”

1.1 一张图看懂内存账本：张量的“生”与“死”

先把一件事说清楚：模型里的每个 tensor，不是从头到尾都活着的。TFLite 拿到一个.tflite文件后，会先构建执行计划（execution plan），把算子排成一个可执行的顺序。在这个顺序里，每个 tensor 都有自己的“出生点”——被哪个算子写出来，以及“死亡点”——最后一次被哪个算子读走。

出生之前，它不占内存；死亡之后，它的内存理论上就可以腾给别人。这段从出生到死亡的时间，我们叫它“存活区间”（lifetime）。内存规划器做的所有事情，本质上都是围绕这张“生死表”来进行的。

我常用的一个类比是自习室座位：同一个座位，上午 A 同学坐，下午 B 同学坐，只要时间不冲突，就能共享。两个张量的存活区间只要没有交集，它们在内存里就可以共用一个位置。这就是整个 arena 分配机制最基本的出发点。

这里要注意一个细节：规划不是“用完立刻释放、下一秒立刻复用”的极致玩法，而是按算子粒度保守计算的。因为一个算子执行时，它要同时持有输入和输出张量，规划器必须保证这个瞬间没有一个张量被提前覆盖。宁可多留一点，也不能冒险早放。

1.2 朴素分配与 arena 分配：一间房一个住客，还是一层楼共享会议室

如果不做规划，最朴素的做法是：每个 tensor 都独立 malloc 一块内存。这就像给每个 tensor 开一间酒店房，互不打扰，也互不复用。代价是总内存等于所有 tensor 体积之和。对一个 224×224 输入的 CNN，中间 feature map 动辄几十层，这个总和会非常吓人，而且大量内存其实只服务同一时刻的少数几个 tensor。

TFLite 现在走的路线是 arena 分配：预先申请一大块连续内存，规划器在这块“竞技场”内部给每个 tensor 分配偏移量（offset），生命周期不重叠的 tensor 共用同一段空间，推理结束后整块一次性释放。

这两条路线的差异，可以用会议室来类比。朴素分配是每个会议租一间独立会议室，无论会议室有没有人，都会一直占着；arena 分配则是整层楼只有几间大会议室，前台根据各会议的起止时间排表，错开的会议共用同一间。楼还是那栋楼，但能容纳的会议数量多得多。

从工程角度看，arena 还有个额外好处：malloc 次数大幅减少，通常整个推理过程只有一次大分配，内存碎片和分配器开销都低。对 Android 这类有 native 内存压力、对 OOM 零容忍的环境来说，这个特性非常关键。

2. 拆开规划器黑盒：生命周期分析、贪心分配与竞技场复用

2.1 三步走：算子扫描、冲突建模、空洞落位

我花了不少时间读 TFLite 的arena_planner.cc，把它干的事归纳成三步：

第一步，扫描 execution plan。对每个算子，遍历它的输入和输出：输入 tensor 更新“最后使用”的算子编号；输出 tensor 记录“第一次写出”的算子编号。模型自身的输入输出有特殊处理，一般会让输入从一开始就存活、输出一直存活到最后，避免执行到一半发生越界。

第二步，建立冲突关系。两个 tensor 的存活区间分别是[first_new, last_used]，当区间存在交集时，它们不能共享内存。这一步会得出哪几个 tensor 可以复用同一块空间。

第三步，贪心落位。把可以复用的一组 tensor 按一定顺序，在 arena 里找空位、对齐、写偏移量。TFLite 的ArenaPlanner干活时拿的是GreedyMemoryPlanner的实现，名字里的 greedy 说明它不是全局搜索最优解，而是用贪心策略快速得到一个足够好的布局。

实际源码里，每个 tensor 在规划完成后都会记下自己的{offset, size}，这个计划会保存在 interpreter 内部。等到真的开始跑模型，arena 内存已经在AllocateTensors()阶段一次性分配好了，后面的Invoke()只是照着计划把数据写到指定偏移位置。

2.2 为什么优先分配大张量：排序、对齐与碎片账

这里有个值得展开的点：到底按什么顺序给 tensor 排座次？

直觉上你可能会按算子执行顺序排，但 TFLite 的做法是把 tensor 按体积从大到小排序，再逐个放入 arena。原因很简单：大的张量最难找空位，先把大块头落位，后面的小块才能见缝插针；反过来小张量先占位，大张量很可能找不到连续空间，arena 里就会留下大量无法使用的空洞。

对齐也是一个容易忽略的成本。arena 基地址和每个偏移量都要满足对齐要求，常见实现按 16 字节或 64 字节对齐，具体随编译目标和平台变化。一个 tensor 申请的大小会被向上取整到对齐边界，多出来的字节属于内部碎片。

举例：某 tensor 实际需要 10242 字节，按 16 字节对齐就得占用 10256 字节，浪费 14 字节。单个看可以忽略，但模型里几百上千个 tensor，加起来也是几十 KB 到几百 KB。好在这种开销通常占 arena 总量的比例很小，属于“知道就好、不用强求”的范畴。

2.3 三种身份的边界：只读权重、可复用激活、持久状态

打开 TFLite 源码的common.h，你会看到张量分配类型有多种，平时最需要关注三种：

• 只读权重（kTfLiteMmapRo）：模型文件里的常量 buffer，理论上可以直接用 mmap 映射，不进 arena、不拷贝、不占 heap。这也是“模型文件 50MB 但运行内存不大”的原理所在。前提是加载方式支持 mmap，如果用自定义方式从流里读取，权重还是会被拷进内存。
• 可复用激活（kTfLiteArenaRw）：绝大多数中间计算结果，生命周期短，是规划器优化和复用的主要对象。
• 持久激活（kTfLiteArenaRwPersistent）：像 LSTM 的隐状态、自定义变量这类需要跨多次推理保留的数据，规划器会把它放在 arena 最前面，预留后不参与复用。

所以你看，规划器并不是万能的。模型里持久状态越多，arena 可复用的空间就越少。这也是循环模型经常比单帧 CNN 更吃内存的一个重要原因。

3. 实操：翻开模型的“内存账本”

3.1 静态盘点：用 Python 把每个张量的体积列出来

纸上谈兵完了，动手。第一步是拿到一个模型后，先做静态盘点。我通常用 TensorFlow 自带的 Python 绑定，几行代码就能把每个 tensor 的形状、类型和理论体积列出来。

import tensorflow as tf import numpy as np model_path = "model.tflite" interpreter = tf.lite.Interpreter(model_path=model_path) interpreter.allocate_tensors() details = interpreter.get_tensor_details() dtype_size_map = { np.float32: 4, np.float16: 2, np.int32: 4, np.int8: 1, np.uint8: 1, np.int64: 8, } rows = [] total = 0 for d in details: # 标量 tensor 的 shape 是空数组，np.prod([]) == 1.0，同样适用 n = int(np.prod(d["shape"])) size = n * dtype_size_map[np.dtype(d["dtype"])] total += size rows.append((d["index"], d["name"], d["shape"], size)) rows.sort(key=lambda r: -r[3]) for idx, name, shape, size in rows: print(f"{idx:4d} {size / 1024:8.1f}KB shape={str(shape):30s} name={name[:60]}") print(f"tensor 总数: {len(rows)}") print(f"理论全量（没有任何共享时的总和）: {total / 1024 / 1024:.2f} MB")

这里算出的“理论全量”是一个上限：假设所有 tensor 各自独立申请内存，谁也不共享。真实运行时会低很多，因为 arena 规划器会把不重叠的 tensor 合并。但先算这个值有个好处——它能告诉你模型结构上的“内存基数”有多大。如果这个基数已经超过设备可用内存，那后面调规划器也没用，得先改模型结构。

如果想进一步区分权重和激活，可以看名字里的关键字，比如weight、kernel、bias、Const通常对应常量。配合 Netron 打开模型，能对照着看更直观。

3.2 动态观测：在 Android/Linux 上读真实峰值内存

静态账本是预测，动态观测才是真相。我一般用两层手段：

第一层，Android 上直接看进程内存。App 跑起来、模型加载完之后执行：

adb shell dumpsys meminfo com.example.myapp

重点看Native Heap那一项，TFLite 的 arena 内存主要落在 native heap 里。多跑几次推理，对比数值变化，能看出是否存在内存持续增长。

第二层，在进程内部打点。我在 C++ 宿主的AllocateTensors()前后分别记录内存占用，差值就是模型加载 + 规划器建 arena 的开销。Linux 上可以读/proc/self/status：

def read_rss(): with open("/proc/self/status") as f: for line in f: if line.startswith("VmRSS"): return line.strip() return "N/A"

注意，VmRSS 是整个进程的物理内存，不只有 TFLite。所以正确用法是在加载模型前记录一次基线值，加载后再记录一次，两次的差值才接近模型引入的开销。

3.3 算一笔账：一个最小 CNN 的 arena 效果

静态分析+动态观测中间的桥，是亲手算一遍生命周期。我拿一个极简 CNN 举例，输入 224×224 RGB 图像，三层卷积后接全连接：

如果每个 tensor 独立分配，总内存是 588 + 784 + 392 + 196 + 4 = 1964KB。现在看存活区间：

• op0 执行时，input 被读取、c1_out 被写入，两者同时存活，峰值 1372KB。
• op1 执行时，c1_out 被读取、c2_out 被写入，峰值 1176KB。
• op2 执行时，c2_out 被读取、c3_out 被写入，峰值 588KB。
• op3 执行时，c3_out 被读取、logits 被写入，峰值 200KB。

真正的峰值是 1372KB，已经比 1964KB 低了不少。而且 c3_out 的存活区间和 c1_out 完全错开，它可以复用 c1_out 在 784KB 偏移上的空间；logits 复用 c2_out 在 392KB 偏移上的空间。最终 arena 大约只需要 1764KB，比朴素方案省约 10%。

这个例子很小，收益看起来不够刺激。换成层数更多、通道数出现波峰波谷的真实 CNN，中间几个大激活互相错开，复用收益会明显放大。但如果网络里有残差连接，某些张量要跨很多算子存活，复用率就会被拉低。这也是为什么残差网络在内存上天然不如纯串联网络“讨喜”。

4. 常见问题与排查技巧实录

4.1 指针过期、内容被覆盖：Invoke 之后别痴迷于中间张量

我第一次在 C++ 里调试 TFLite 时踩过一个很经典的坑。我在某次Invoke()之后保存了一个中间 tensor 的data指针，下一帧再读取，发现内容变成了一堆垃圾值。

原因不复杂：arena 规划器判定这个张量已经死亡，把它的内存空间分配给了另一个张量，内容被覆盖了。这在 arena 机制里是正常行为，甚至是设计目的。所以如果你要保留中间层结果做调试或可视化，必须及时拷贝出来，或者打开调试开关。

较新版本的 Python 绑定里，tf.lite.Interpreter有experimental_preserve_all_tensors参数，打开后推理不会释放中间张量，方便调试。但要注意，这只是调试选项，会显著增加内存，别带着它上生产。

4.2 动态 shape 导致反复重规划，内存不降反升

做视频或音频流处理时，我见过有人每帧都调用resize_tensor_input()改输入尺寸。TFLite 会在下一次AllocateTensors()时重新执行整个规划流程。重规划本身的计算量不大，但它可能会释放旧 arena、申请一块更大的新 arena，在申请新内存时旧内存还没释放，RSS 瞬时可能涨到两倍。

我的建议是：输入尺寸能固定就固定。把输入统一 resize 到模型支持的最大尺寸，在预处理阶段做缩放，而不是去改 interpreter 的输入 shape。batch 也固定为 1，端侧推理场景几乎不需要动态 batch。

还有个隐蔽的重规划来源：模型内部存在动态输出形状的算子（比如动态 shape 的 Resize），导致某个 tensor 被标记为动态分配（kTfLiteDynamic）。这类 tensor 不走 arena，而是每次执行时临时 malloc，既破坏复用，又容易造成碎片和峰值上涨。排查时可以在 C++ 里检查tensor->allocation_type，发现动态 tensor 就该考虑换算子实现或固定 shape。

4.3 int8 量化后内存账怎么重算

量化是降低内存最直接的手段，没有之一。int8 量化后，激活从 float 的 4 字节变成 1 字节，权重同理；float16 量化则是权重减半。你可以用 3.1 节的脚本重跑一遍，把dtype映射改成 int8/uint8，naive 理论值立刻能看出差距。

但注意两个坑。第一，部分量化只量化了支持 int8 的算子，其余算子仍以 float 运行，混合模型里 float 中间张量照样占 4 字节。第二，int8 激活虽然占 1 字节，但很多算子的中间累加器（accumulator）还是 int32，这部分空间规划器也会一并管理。

团队里有人问我“量化后是不是就不用看内存了”，我的经验是恰恰相反——量化后更需要看。因为当你把 4 字节压到 1 字节后，arena 碎片和对齐浪费占比会相对变高，plan 的细节差异开始影响最终结果。

4.4 委托（Delegate）把内存搬走了，但你没注意搬运工

用 GPU delegate 或 NNAPI 时，部分算子会被换到加速器上执行。这时候很多人的直觉是“CPU 内存能省一大截”，但实际不一定。TFLite 仍然要为这些算子保留 CPU 侧的输入输出缓冲，数据在 CPU 和加速器之间搬运时，两边经常要各放一份。小模型委托之后，内存可能不降反升。

所以判断 delegate 是否值得，不能只看推理速度，要同时看 RSS。做法是分别跑一遍纯 CPU 和 delegate 两种配置，用 3.2 节的手段对比 native heap 峰值。一般情况下，我建议先量化、再考虑委托；量化往往比委托更能省内存，而且没有数据拷贝的额外负担。

5. 我的实操体会与几条建议

5.1 优化顺序：先定 shape，再量化，最后才谈规划器

我把一套固定套路用在了多个项目里：先固定输入尺寸和 batch，再用 3.1 节的脚本算 naive 总量，接着做 int8 量化重算，最后才在真实设备上测 RSS。如果你一开始就纠结 arena 的对齐细节，方向大概率错了。

这个顺序背后有个很朴素的道理：内存优化的杠杆大小不同。改 shape 能按面积倍数压缩激活，量化按 4 倍压缩，而 arena 规划器只能帮你在既定结构里挤出复用空间。杠杆大的手段优先，剩下的交给管家处理。

5.2 把 RSS 曲线当成优化仪表盘

每次改动模型或推理配置，我都会维护一张简单的 RSS 对比表：模型加载前、加载后、做完一次推理、连续推理 100 次这四组数值。连续推理多次后内存还在爬升，基本可以断定有动态 tensor 或 delegate 缓冲没释放干净。RSS 是比任何理论估算都诚实的裁判。

另外一个小技巧：把get_tensor_details()的结果导成 CSV，按体积排序后存档。每做一次量化或结构改动，对比 top10 大 tensor 的变化，你能直观看到优化到底砍在哪一层，也方便向同事解释“内存为什么降下来了”。

5.3 什么时候才需要自己写规划器

最后给一个诚实的结论：绝大多数项目根本不需要自己写内存规划器。TFLite 内置的 arena 机制已经覆盖了标准 CNN、Transformer、LSTM 等主流结构的复用场景。真正需要自定义 allocator 的场景，通常是跨模型共享同一块缓冲、内存预算有硬上限，或者裁剪到极致的 MCU 环境。

对我个人来说，理解这套机制最大的价值，不是去改写源代码，而是在内存出问题时能快速判断：是模型结构的问题、量化的问题、动态 shape 的问题，还是 delegate 的问题。这些问题都排除后，你甚至可以不再打开arena_planner.cc，因为管子已经帮你把一切都安排好了。