尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

TFLite内存规划器:模型小却内存暴涨的根源与优化

TFLite内存规划器:模型小却内存暴涨的根源与优化 前阵子我排查一个 TFLite 检测模型的内存暴涨问题模型权重才 4MB输入也就 640x640理论上中间激活峰值撑死 5MB结果在手机上跑起来 RSS 直接飙到几十 MB。顺着代码往下挖发现很多人把推理引擎当黑盒在用完全不关心里面的内存规划器。其实 TFLite 这类推理引擎里内存管理从来不是“用多少申请多少”那么简单的它靠一个藏在底层的“内存管家”统一规划——这就是标题里说的 TFLite 内存规划器。这篇内容适合正在做移动端或端侧部署的工程师也适合那些想搞清楚推理引擎内部实现的学习者。我会从“它到底解决什么问题”开始讲清楚张量生命周期、Arena 分配的核心思路再带你手动推演一遍分配过程最后给出一套排查和调优的实操方法。读完你至少能回答三个问题为什么模型文件小但运行时内存大、中间张量为什么能互相复用内存、以及当 arena 暴涨时该从哪里下手。1. 推理引擎里的“隐形内存管家”1.1 内存压力到底从哪来很多人有一个直觉模型文件 4MB那推理时内存应该也就 4MB 出头。这个直觉在 PC 上误差不大但在移动端完全不是一回事。模型文件里最多的是权重属于只读数据而推理过程中每一层都会产生中间张量比如卷积的输出、激活函数的输出、每次 Pooling 后的特征图。一个 100 层的网络逐层产生、持有、释放这些张量如果没有统一规划内存会变成所有层输出张量的大杂烩。举个例子一个普通的 MobileNetV2输入一张 224x224 的图中间特征图尺寸千变万化有的是 112x112x32有的是 14x14x1280。如果每个算子都独立 malloc 一块内存那么总内存就是这些张量大小的简单叠加峰值很容易到几十 MB。但仔细观察会发现很多中间张量的生命周期根本不重叠——前一层算完它的输出不会被后面所有层一直用等某个节点执行完这块内存其实就“死”了。推理引擎的调度单位是算子节点所有张量在计算图中都有自己的“出生点”和“死亡点”。内存规划器要做的就是分析这些生命区间把那些“死掉”又没人占用的内存腾出来给后面的张量继续用。说白了这就是一个“会上引座员”的角色你走了座位马上给下一批观众。1.2 内存规划器具体管什么事TFLite 把这件事交给了专门的模块在源码里对应的是ArenaPlanner和SimpleMemoryArena。你可以把ArenaPlanner理解为“策略制定者”它负责在真正执行模型之前把所有张量的内存位置预先算好SimpleMemoryArena则是“仓库管理员”它维护一块连续的大内存按规划器给出的偏移地址把每个张量放进对应位置。这里的关键词是“预先算好”。推理引擎不想在运行过程中频繁走 malloc / free 流程那样既慢又容易产生碎片。TFLite 的做法是在AllocateTensors()阶段做一次性规划先分析计算图给每个中间张量分配一个相对 arena 首地址的偏移量然后申请一大块连续内存运行时所有张量都在这块内存里搬进搬出。执行一个算子时只需要通过偏移量定位输入输出张量即可完全没有动态分配的负担。还有个容易被忽略的点内存规划器不只是管“中间激活”它还要区分不同类型的张量。比如模型的权重TFLite 可以通过 mmap 直接映射模型文件不占用额外内存而 LSTM 内部的 state 是持久变量跨时间步、跨推理调用都要保留不能和普通中间张量一样用完就回收。这些分类决策也由规划器在AllocateTensors()时统一处理。1.3 没有规划器会怎样我在自研一个小型推理引擎的时候试过偷懒每个算子执行完就把输出动态分配出去下一个算子用完再释放。结果就是内存碎片严重、分配次数多到影响帧率更麻烦的是内存不确定性——某次输入尺寸稍微变一下峰值内存就会波动甚至出现 OOM。嵌入式环境里 malloc 本身就有风险频繁调用会让堆碎片越来越严重跑一会儿就容易内存不足。规划器的好处是“离线决策、在线执行”。所有内存决策在初始化阶段就定死运行阶段不产生分配行为时间和空间都是可预期的。这也是为什么 TFLite 在移动端能保持稳定的内存表现而不是像某些通用框架那样动不动就浮点数秒级推理、内存却下不来。2. 张量生命周期与可复用内存2.1 看清 TFLite 的张量分类理解内存规划器之前得先认识 TFLite 里的张量分类。每个张量都有一个allocation_type字段大致分成几类kTfLiteMmapRo只读映射通常是权重、偏置直接从模型文件映射到内存不参与中间内存规划。kTfLiteArenaRw普通的可读写中间张量生命周期完全在计算图内部是 Arena 规划的核心对象。kTfLiteArenaRwPersistent可读写但需要在多次推理或子图调用之间保留的张量比如控制流里的循环变量、LSTM 的 state。kTfLiteDynamic动态张量运行时尺寸可能变化无法提前规划只能单独动态分配。kTfLiteMemNone还没分配内存的张量通常是一个占位。从内存规划角度kTfLiteArenaRw是最值得优化的对象。它只活在计算图内部生命周期完全由节点顺序决定kTfLiteArenaRwPersistent则因为要跨调用存活会被放到另一块持久内存区域不能被普通中间张量复用而kTfLiteDynamic是规划器的“盲区”无法离线决定位置。实际调试时我会在AllocateTensors()之后扫一遍所有 tensor看它们的allocation_type分布。如果发现大量中间张量被标成了kTfLiteDynamic那就得警惕了——这些张量内存是动态分配的arena 想管也管不着。2.2 张量的“出生”与“死亡”内存规划器最核心的输入是每个张量的生命周期。生命周期怎么定义很简单计算图里的每个算子节点都有一个编号按执行顺序排张量被某个节点“生产”出来又被某些节点“消费”。它的出生点是生产它的节点编号死亡点是最后一个消费它的节点编号。用真实代码里的逻辑来表达就是遍历每个节点更新它的输入和输出张量的first_use与last_use// 生命周期收集伪代码概念来自 TFLite ArenaPlanner 的 prepare 阶段 for (int node_id 0; node_id graph.nodes_size; node_id) { for (int input : graph.nodes[node_id].inputs) { if (first_use[input] -1) first_use[input] node_id; last_use[input] node_id; } for (int output : graph.nodes[node_id].outputs) { if (first_use[output] -1) first_use[output] node_id; last_use[output] node_id; } }注意一个张量如果是某个节点的输入又恰好是另一个节点的输出那么它在这个节点上相当于“被消费后再被生产”。很多内存规划器会把这种情况当成寿命连续处理但 TFLite 里有些算子支持原地更新in-place也就是输出直接复用输入的地址此时生命周期标记要格外小心。一旦有了first_use和last_use每个张量就变成一条时间轴上的线段[first_use, last_use]。这条线段是内存复用的基础。两个线段如果完全不重叠它们对应的张量就可以共享同一块物理内存。这就是张量内存复用最朴素也最有效的判据。2.3 什么条件下两个张量能共享内存两个张量要共享内存需要同时满足三个条件第一生命周期不相交。张量 A 的last_use必须早于张量 B 的first_use或者反过来。也就是说没有任何一个节点需要同时用到 A 和 B。只要有一个算子既要读 A 又要写 B它们就只能分占两块内存。第二尺寸匹配。共享内存时物理块的大小是两者中的较大值。假设 A 是 100 字节B 是 50 字节A 释放后 B 可以放进 A 的位置但反过来不行。这让分配器需要维护一个“可用空间列表”按大小查找合适的内存块。第三对齐要求。TFLite 的 SIMD 算子比如 ARM NEON对内存地址有对齐要求默认对齐通常是 64 字节。即使生命周期和尺寸都满足偏移地址也要对齐到 64 字节边界否则算子访存可能会触发 bus error 或性能回退。拿生活里的例子类比会议室要预约同一个会议室可以给下午和晚上两拨人开会用只要时间不撞车但如果你要开一个要坐 100 人的大会就得保证会议室容量足够而且投影仪位置还得合适。内存复用就是这个道理不是简单地把地址倒来倒去而是要统筹时间、大小、对齐三个维度。3. ArenaPlanner 的规划流程3.1 第一步建立全图张量生命表真正构建生命表的代码要比上面的伪代码复杂一些因为要考虑节点内部多个输入输出之间的关系。规划器会把每个临时张量包装成一个TensorUsage之类的结构记录张量编号、尺寸、首次使用节点、最终使用节点、是否持久化等信息。TFLite 里有一点很关键它规划的是“算子执行顺序”下的生命周期而这个顺序在 FlatBuffer 模型里基本是固定的转换器生成模型时已经确定了节点的先后关系。也就是说同样一份网络结构因为算子存储顺序不同张量生命周期就可能不同最终 arena 大小也会不一样。这是很多人不知道的细节我后面会再展开。构建生命表时还有一个隐含步骤确认哪些张量真正需要参与规划。输入张量、输出张量通常由外部提供不需要在 arena 内分配权重张量是MmapRo也不参与只有那些ArenaRw和ArenaRwPersistent类型的中间张量才会被纳入分配池。这一步看起来简单但一旦模型里混入了控制流子图张量的归属关系会变得复杂需要在生命表里显式标记。3.2 第二步贪心打包与内存分配有了生命表之后ArenaPlanner 开始按“出生顺序”或“请求顺序”逐个处理中间张量。每一步做两件事先回收已经死亡张量占用的内存块再为当前张量找一个合适的空闲块。这里的贪心策略是只要有空闲块且大小足够就用没有就去 arena 尾部追加。为什么用贪心而不是全局最优因为内存分配问题本质是一个装箱问题bin packing要求最优解是 NP-hard 的。TFLite 面对的是各种结构完全不同的模型规划器必须在毫秒级完成计算不能跑一个整数规划求解器。实践也证明贪心算法在神经网络这种生命周期比较规整的图上已经能接近最优解损失一般很小。回收逻辑有个细节必须提醒你不是在“新张量申请的节点开始时”统一释放而是要精确到“上一个算子执行结束之后”。一个张量可能在节点 5 被最后一个算子消费那么节点 5 执行完才能释放。但如果新张量也是节点 5 的产物那么它和旧张量严格来说是在同一时刻完成交接只要算子本身安全它们可以复用同一地址。TFLite 的分配器允许这种“在同一节点边界交接”的复用这也是很多手工分配时容易疏忽的地方。3.3 手动推演一个小图怎么分配内存下面我用一个简化模型演示完整的偏移计算过程。假设有四个算子节点 0、1、2、3每个节点产生一个中间张量张量的大小和生命周期如下张量尺寸字节生命周期节点区间t164[0, 1]t2128[1, 2]t364[2, 3]按贪心分配同时要求 64 字节对齐第一步处理 t1。它出生在节点 0直接放到 arena 头部偏移 0占用 [0, 64)。此时 arena 尾部是 64。第二步处理 t2。t2 出生在节点 1但 t1 要活到节点 1 结束所以在 t2 出生时 t1 还占着 [0, 64)。不能用这块空间只能放到 arena 尾部。因为 64 字节已经对齐t2 偏移取 64占用 [64, 192)。arena 尾部变成 192。第三步处理 t3。t3 出生在节点 2此时 t1 已经在节点 1 结束后死亡。t3 的尺寸是 64刚好可以放进 t1 留下的 [0, 64) 空洞。于是 t3 被放到偏移 0arena 尾部仍然只有 192。最终结果arena 总大小 192 字节。如果不用内存规划各自独立分配那需要 6412864256 字节。省下来的 64 字节来自 t1 和 t3 的生命周期复用。真实模型里这种复用的效果会被放大很多倍因为中间张量动辄几百上千个复用率往往能达到 50% 以上。3.4 关键数据结构与分配伪代码ArenaPlanner里主要维护两类东西一个是“活跃张量”集合记录当前还活着的张量和它们占用的内存块另一个是“空闲块”集合记录已经释放、可以被新张量复用的空间。下面是分配阶段的概念性伪代码基本反映了 TFLite 的设计思路struct Block { size_t offset; size_t size; }; // active_blocks: 当前活跃张量 - 占用块 // free_blocks: 可用块列表按 offset 排序 // tail: arena 已使用的尾部 void AllocateForNode(int node_id) { // 1. 先回收已经结束生命的张量 for (auto tensor : active_tensors) { if (lifetime[tensor.id].last_use node_id) { auto block active_blocks[tensor.id]; free_blocks.push_back(block); active_tensors.erase(tensor); } } // 2. 分配当前节点的输出张量 for (int tensor_id : node.outputs) { if (allocation_type[tensor_id] ! kTfLiteArenaRw) continue; size_t size align_up(tensor_size[tensor_id], kAlignment); // 3. 找一块足够大的空闲块否则追加到尾部 auto it std::find_if(free_blocks.begin(), free_blocks.end(), [size](const Block b) { return b.size size; }); if (it ! free_blocks.end()) { active_blocks[tensor_id] {it-offset, size}; free_blocks.erase(it); } else { active_blocks[tensor_id] {tail, size}; tail size; } active_tensors.push_back(tensor_id); } }为什么维护free_blocks而不是只记住“上一个结尾”因为张量释放的顺序和申请顺序不一样会产生碎片空洞。比如大块 A 释放后小块 B 申请可能放进 A 的位置但之后大块 C 申请又因为放不进 A 的残留而继续涨尾部。分配器通过维护空闲块列表来尽量利用空洞。TFLite 还做了一点优化优先使用 offset 最小的空闲块让尾部增长尽量慢这样 arena 总大小更紧凑。3.5 对齐问题看起来小踩到就崩TFLite 默认对齐是 64 字节这个值的来历主要照顾 ARM NEON 的向量访存需求。对齐不仅影响性能还影响正确性。某些内核会直接加载 16 字节或者 32 字节的向量如果内存地址没对齐轻则速度暴跌重则直接崩溃。计算偏移时实际使用的是align_up(offset size, 64)之后的值。也就是说每个张量块之后可能留下几十字节的碎片。不要小看这点浪费如果模型有几百个中间张量对齐碎片可能累计到几十 KB。但相对整体 arena 来说这点代价是值得的。调试时如果只在 x86 机器上跑可能察觉不到对齐问题因为 x86 对未对齐访问容忍度更高。一旦部署到 ARM 手机或者 Cortex-M 系列芯片上问题就会暴露。所以我建议在内存规划器代码里显式打印每个张量的 offset 和 size检查它们是否都满足 64 字节对齐这个习惯能省去很多底层排查时间。4. 实操如何拿到 TFLite 的真实内存画像4.1 用代码直接读 arena 信息最理想的自然是调用公开 API 拿到内存统计。TFLite 在不同版本里提供的接口不完全一致老版本在Interpreter上可以直接访问内部 arena新版本把实现藏得更深了。一个稳妥的办法是在源码构建时给ArenaPlanner加日志或者在调用AllocateTensors()之后遍历所有张量查看它们的地址和大小。给一个思路层面的参考代码具体 API 以你用的源码版本为准// 伪代码AllocateTensors 完成后统计每个 tensor 的内存信息 auto* interpreter /* 你的 TFLite Interpreter 实例 */; interpreter-AllocateTensors(); size_t arena_total 0; for (int i 0; i interpreter-tensors_size(); i) { TfLiteTensor* t interpreter-tensor(i); if (t-allocation_type kTfLiteArenaRw || t-allocation_type kTfLiteArenaRwPersistent) { arena_total align_up(t-bytes, 64); Printf(tensor %d: offset%p, size%zu, type%d\n, i, t-data.data, t-bytes, t-allocation_type); } }这里有个小技巧打印t-data.data指针把所有中间张量的地址拿出来画成区间图。如果发现两个生命周期不相邻的张量地址一样说明规划器成功复用了内存如果地址完全分散且没有重叠说明可能没有启用规划器或者模型里大部分是动态张量。如果你的设备环境不方便跑 C也可以在 Python 侧通过tf.lite.Interpreter的get_tensor_details()拿到每个 tensor 的 shape、dtype但拿不到精确的 arena offset。这种情况下我更推荐走 4.2 节的工具路线。4.2 用 Benchmark 工具和系统命令交叉验证TFLite 官方提供了一个 benchmark 工具在tensorflow/lite/tools/benchmark目录下支持把内存统计打到日志里。你可以在初始化阶段传入一个内存监控器在AllocateTensors前后记录 malloc 统计。虽然具体参数每个版本有调整但思路大同小异跑一个空输入循环记录峰值 RSS再单独把 arena 大小算出来对比。在 Android 设备上我还会配合系统工具一起看。比如adb shell dumpsys meminfo package_name重点看 Native Heap 和 Graphics 内存。TFLite 的 arena 属于 Native Heap如果 arena 规划得好Native Heap 曲线应该很平稳不会随着单次推理次数增长而持续上升。如果每隔几次推理 Native Heap 就涨几 MB那大概率是存在动态张量没有走 arena每次都触发新的分配。另一个有用的方式是统计“单次推理触发的 malloc 次数”。在 TFLite 的执行循环里打桩记录每个算子执行时的分配调用。正常情况下AllocateTensors 之后运行阶段不应该有任何 malloc。如果发现运行阶段还有分配就顺着调用栈找到底是哪个算子在做动态内存申请通常是自定义算子或者ResizeInputTensor导致的重规划。4.3 改源码加日志看每个张量的规划结果最快的方案其实就是给arena_planner.cc加打印。你不需要理解全部算法只要在PlanAllocations结束后遍历分配的记录把 tensor 编号、偏移、大小打印出来。日志格式可以这样 ArenaAlloc: tensor12 size12800 offset512 ArenaAlloc: tensor37 size25600 offset0 TotalArenaSize: 1048576把这些输出保存下来用脚本画成泳道图——每个 tensor 一条横线按生命周期排在时间轴上纵轴表示 offset。一眼就能看出哪些区域被重复使用、哪些大块张量拉高了整体峰值。这个手段比任何高级 profiler 都直观而且不依赖外部工具。如果你不想修改官方源码可以复制一份ArenaPlanner改成自定义 planner 挂进 interpreter在自定义版本里加各种统计。TFLite 的InterpreterBuilder允许你传入自己的内存规划器代价是要重新实现接口。对只是想排查问题的团队我觉得临时加日志更划算。5. 常见问题与调优实录5.1 模型文件不大arena 却暴涨这种问题最常见的原因就是中间张量太多或者生命周期重叠太严重。多分支结构的网络比如 FPN、Attention 特征融合会同时存活多个大张量这些分支的输出都要保留到融合节点导致同一时间活跃张量非常多。排查步骤我一般这样走先统计张量数量再统计平均生命周期长度。如果一个图里有几百个中间张量平均生命周期超过 10 个节点那说明很多张量没有被及时释放arena 被拉长是必然的。压缩手段优先级从高到低排列第一量化。从 float32 降到 int8中间张量直接缩小四倍这是最暴力的降内存手段。第二激活函数在算子内部原地更新避免产生新的张量。第三手工融合算子比如 ConvBNReLU 融合成单个算子减少中间张量数量和生命周期。第四在保证精度的前提下简化网络结构减少分支宽度或特征图分辨率。这里要特别提醒量化之后很多中间张量仍然是 int8但某些算子为了精度会暂时把数据转成 float32 再转回 int8这种临时 buffer 不算在图张量里却会额外占用内存。排查时不要只盯着 arena还要关注算子内部申请的临时 buffer。5.2 动态尺寸张量带来的麻烦TFLite 对动态 shape 支持得比较弱主要原因是内存规划器必须在运行前把所有位置定死。如果某个张量运行时尺寸变了原来的偏移和大小就失效了要么重新规划要么走动态分配。重新规划的成本很高而且会导致 arena 大小在新尺寸上重新计算之前的紧凑布局全部作废。我在项目里遇到过一个检测模型输入尺寸有时候 320x320有时候 640x640。TFLite 在每次ResizeInputTensor之后都要重新AllocateTensorsarena 按最大尺寸规划。问题是模型里如果有些分支依赖于输入尺寸张量大小会成倍增长arena 峰值按最坏情况预留内存一下子膨胀好几倍。解决思路有两个方向一是固定输入尺寸宁愿训练时做多尺度也不要部署时随意改二是把动态张量尽量控制在算子内部比如用自定义算子在内部按需分配不要让它成为计算图张量。如果自定义算子内部确实需要动态内存尽量复用已有 buffer像TfLiteContext提供的GetScratchBuffer那样避免每次推理都重新分配。5.3 同一个地址被两个张量复用后结果出错这是个很隐蔽的坑。你启用了 arena内存规划得很好但程序变随机出错了。原因往往是某个算子在同一时刻引用了两个已经被规划到同一个偏移上的张量。正常情况下规划器不会让活跃区间重叠的张量复用同一地址但边界情况处理不好就会出问题。一个常见场景是 in-place 算子。比如某个算子声明自己可以原地执行输出直接写到输入上那么它的输出张量和输入张量可以共用一个偏移。如果这个算子其实并不完全支持原地操作或者实现里有隐藏的逻辑依赖原始输入复用就会导致数据被提前覆盖。排查方法很简单把输入张量和输出张量强制放到不同偏移再看问题是否消失。如果不再出错就是 in-place 声明和实现不一致。另一个场景是张量生命周期边界计算有误当一个张量是节点 N 的输入另一个张量是节点 N 的输出如果生命表把前者标记为从 N 结束后才死亡、后者标记为 N 开始时就出生它们就可能错误重叠。遇到这类问题我会在自定义算子上下文里禁用kAllowInPlace之类的优化开关优先保证正确性。等内存优化做完之后再认真审查每个算子的原地操作实现逐步打开复用。5.4 算子顺序对峰值内存的影响有多大这个点不少人忽略。同一个计算图可以把几个无关算子排在不同顺序只要满足拓扑序就行。但算子顺序会直接改变张量的死亡时间从而改变生命周期重叠程度。举个简单的例子一个网络先做两个独立分支每个分支再各出一个结果最后融合。如果两个分支的算子交叉执行其中一个分支的中间张量一直存活另一个分支的中间张量也在存活峰值就高。如果先把一个分支完整算完再算另一个前一个分支的中间张量早早就释放了arena 可能更小。TFLite 转换器大多时候会按模型原始定义的节点顺序生成执行顺序但这个顺序不一定是最优的。你在转换前后对比 arena 大小时如果发现同样结构不同版本的 TFLite 模型占用内存不同很大原因就是算子顺序变了。想优化的话可以考虑在模型图优化阶段做节点重排把生命周期长的张量尽量往后放让它们少跟其它大张量重叠。这个优化对内存的影响不输于量化但实现起来需要动图优化代码。5.5 多实例和子图之间如何共享内存如果你的进程里同时跑多个 TFLite interpreter 实例比如一个检测模型加一个分类模型它们的 arena 是各自独立的内存不会自动复用。想共享内存的话要么把两个模型合并成一个大的计算图要么自己做内存池把两个 arena 放在同一个 buffer 里分别计算偏移再错开分配。子图的情况相对复杂。TFLite 支持控制流算子比如 while 循环子图会持有自己的张量集合。从执行角度看子图和主图交替运行如果每个子图都单独规划 arena内存可能会有重复但它们不是同时存活理论上可以复用。具体实现里子图的持久张量对接主图的持久 arena而内部临时张量通常不跨子图共享所以多子图模型的内存规划是偏保守的。如果你在跑语音或流式模型还会遇到跨推理调用的状态张量。这类张量必须存在持久区不能走普通复用逻辑。规划器把ArenaRwPersistent单独管理避免和一个推理内的临时张量混在一起被错误释放。对这类模型我建议把状态张量显式连接成计算图的输入输出让框架能识别它们跨调用存活而不是依赖黑魔法。6. 工程上的一些建议6.1 把 arena 大小当作重要性能指标我在部署流程里会把“arena 总大小”和“单次推理 malloc 次数”两个指标打进 CI 日志每次更新模型或框架库时对比。模型转换后AllocateTensors()会输出一个 arena 大小如果一次模型更新让 arena 暴涨 20%即使精度没变我也会警觉去查张量生命周期是不是被打散了。这个习惯帮我避免过好几次线上 OOM。很多团队只盯模型文件和推理耗时忽视了内存峰值结果一上真机就闪退。你不需要理解每一个底层细节只要把 arena 大小纳入版本对比体系就能提前发现问题。TFLite 源码里我印象最深的注释是关于为什么用 arena 而不是直接用 malloc保持执行路径的确定性。这个思想我后来移植到了自己的引擎里每次模型初始化时做一次离线规划执行阶段零分配。看起来只是节省了一点 malloc 开销但换来的是帧率曲线的平滑和内存的可预测性这在实时音视频处理场景里非常值钱。6.2 让算子尽量支持 in-place 更新内存规划器虽然可以把生命周期不重叠的张量复用到同一块内存但如果能让“输出直接覆盖输入”那连生命周期重叠都不用担心了。像 ReLU、某些归一化算子它的输入在计算完后就不再需要了完全可以在原地改内存。实现的时候要注意in-place 不是简单的“输入输出索引相同”还要考虑算子内部有没有额外的中间 buffer。比如某个算子需要把输入先搬一份再做计算那它就不适合原地操作。我一般先让普通算子保证正确性再逐个验证是否可以安全开启 in-place。这个收益在通常的 CNN 里可能只有百分之几但在一些连续层很多的小模型上能明显降低峰值。6.3 自己写引擎时怎么参考这个设计如果你也在写自己的推理引擎不用照抄 TFLite 的代码但可以抄它的设计思路把“内存决策”和“执行”彻底分离。初始化阶段收集张量生命周期、计算偏移运行阶段完全按预设地址执行。这样能保证执行路径没有动态分配也方便做内存峰值预测。实现时先从最朴素的贪心开始不要一开始就想实现最优装箱算法。贪心在大多数图上表现足够好而且实现简单调试方便。等到内存成为瓶颈时再考虑更复杂的策略比如按生命周期区间图做最大重叠消除或者引入重计算换内存的思路——那就完全是另一个话题了。最后分享一个我常用的“土办法”在加载模型后把 arena 里每个张量的起始地址和大小输出然后手工挑几个关键层对比。有一次我发现某个 14x14x512 的特征图被分配到了 arena 的开头而 56x56x256 的特征图被分配到了末尾结果中间空闲了一大段。原因就是生命周期算法把一个大尺寸张量排在了前面小尺寸张量无法填满它留下的空洞。后来我把大尺寸张量优先分配arena 又压下去一截。这种细节工具文档里不会写但实际排查时帮助极大。希望这篇东西能帮你少走点弯路也算是我跟内存规划器“过招”这么久的一点经验积累。
返回列表