全解析:从 Tensor Arena 到三阶段分配流程)
人工智能深度学习推理引擎本地部署嵌入式物联网【免费下载链接】tflite-microInfrastructure to enable deployment of ML models to low-power resource-constrained embedded targets (including microcontrollers and digital signal processors).项目地址https://gitcode.com/gh_mirrors/tf/tflite-micro点击查看免费下载本文基于 tensorflow/lite/micro/docs/online_memory_allocation_overview.md 展开结合 TFLM 仓库中MicroAllocator、GreedyMemoryPlanner、SingleArenaBufferAllocator等核心实现与测试代码系统讲解 TFLM 如何在资源受限的嵌入式设备上通过单一 Tensor Arena 完成模型的在线内存规划。读完本文你将掌握 TFLM 的 Arena 头/尾/临时三段式内存布局、flatbuffer 内联缓冲的复用机制以及 Model Init / Prepare / Finish 三个阶段中每一次分配的底层调用链与关键数据结构从而能准确估算模型内存占用、定位分配失败问题。什么是在线内存分配Online Memory AllocationTensorFlow Lite MicroTFLM的目标是让机器学习模型跑在微控制器、DSP 等低功耗、资源受限的嵌入式目标上。这类平台上通常没有动态内存分配器malloc/free或即便存在也不允许在推理热路径中使用。因此 TFLM 采用**在线内存规划online memory planning**策略在运行时读取 flatbuffer 中的模型描述把推理所需的全部中间张量、算子运行时数据、scratch buffer 一次性规划并放置到一个预先分配的连续字节数组中。在线对应的是离线offline离线方案在宿主机PC上预先计算好每个张量的偏移量并写入模型元数据运行时按图索骥而在线方案则由MicroAllocator在AllocateTensors()或首次Invoke()时实时完成布局。本文的主角就是这条在线路径。单一 Tensor Arena一切分配的主战场在线内存规划把所有的分配都策略性地放入一个单一的uint8_t缓冲区数组——即Tensor Arena。这个缓冲区被划分为两个主要区段Head头部放置非持久non-persistent分配例如算子执行时共享的中间张量缓冲区Tail尾部放置持久persistent分配例如TfLiteEvalTensor结构体数组、NodeAndRegistration结构体数组、算子的运行时量化数据等生命周期与整个应用一致。在 Head 与 Tail 之间还有一段Temporary临时区域用于作用域式的短期分配详见 memory_management.md 中的完整三段式图解。Arena 可以由调用方直接传给tflite::MicroInterpreter构造函数也可以先创建tflite::MicroAllocator实例再传入// Buffer for the tensor arena: size_t tensor_arena_size 2048; uint8_t tensor_arena[tensor_arena_size]; // Interpreter using the shared tensor arena above: tflite::MicroInterpreter interpreter( tflite::GetModel(my_model_data), ops_resolver, tensor_arena, tensor_arena_size);从源码层面看Arena 的实际分配操作由SingleArenaBufferAllocator完成single_arena_buffer_allocator.h。它同时实现INonPersistentBufferAllocator与IPersistentBufferAllocator两个接口AllocatePersistentBuffer()从 Tail高地址、向下增长分配AllocateTemp()从 Head 末端低地址、向上增长分配临时缓冲且不会更新 Head 的真实水位直到ResetTempAllocations()被调用。MicroAllocator的类注释中给出了精确的内存布局示意micro_allocator.h************** .memory_allocator-GetBuffer() Tensors/Scratch buffers (head) ************** .head_watermark unused memory ************** .memory_allocator-GetBuffer() -GetMaxBufferSize() - -GetDataSize() persistent area (tail) ************** .memory_allocator-GetBuffer() -GetMaxBufferSize()头部的内存规划器Head 区段的长度并非随意设定而是由tflite::GreedyMemoryPlanner管理greedy_memory_planner.h。它分析整个模型图尽可能复用缓冲区从而让 Head 长度最小化。其核心贪心策略从源码注释可以确认客户端通过AddBuffer(size, first_time_used, last_time_used)登记每个缓冲区及其生命周期首次查询偏移量时触发CalculateOffsetsIfNeeded()计算布局所有缓冲区按大小降序排列最大的缓冲区放在偏移 0依次遍历其余缓冲区在同时活跃的缓冲区之间寻找第一个能容纳当前缓冲区大小的空隙若找不到足够大的空隙则放在最后一个同时活跃缓冲区之后。源码同时诚实声明这不是最优解最优布局是 NP-Complete 问题但在实践中足以产出不错的布局。此外GreedyMemoryPlanner::preserves_all_tensors()返回false意味着推理过程中未被使用的张量数据可能被覆盖——这是共享内存规划的固有特性。每个缓冲区约需 36 字节的规划 scratch 内存per_buffer_size()。flatbuffer 中已有的缓冲区直接复用不重复拷贝TFLite flatbuffer 模型里携带了运行模型所需的各类信息。在线内存规划器会遍历主 subgraph找出模型所需的全部张量——运行时以TfLiteTensor和TfLiteEvalTensor两种 C 结构体表示。其中持久张量如权重张量在 flatbuffer 中是内联inline存放的即权重数据本来就打包在模型文件里。在线规划时这些缓冲区会被直接复用对应的 C 结构体TfLiteEvalTensor的data指针会指回 flatbuffer 中打包好的缓冲区而不在 Arena 中再分配一份拷贝。这意味着权重等常量数据不会占用 Tensor Arena 空间节省了嵌入式设备上极为宝贵的 SRAMArena 只承载可变中间张量、scratch buffer 与各类运行时结构体。三阶段分配流程总览在线模型分配由MicroInterpreter驱动整个流程分为三个阶段阶段触发方式核心方法主要产出Model Init首次Invoke()或显式AllocateTensors()MicroAllocator::StartModelAllocation()SubgraphAllocations张量/节点结构体数组Model PrepareInit 完成后逐算子回调各 kernel 的TfLiteRegistration::prepare()FinishPrepareNodeAllocations()校验、scratch buffer 请求、量化运行时数据Finish Model Allocation所有算子 prepare 完毕MicroAllocator::FinishModelAllocation()提交静态内存计划、固定 Head/Tail 布局MicroInterpreter的公开接口在 micro_interpreter.h 中定义AllocateTensors()显式触发分配Invoke()则在内部自动完成首次分配。两个构造函数分别支持直接给 Arena和传入已创建的MicroAllocator两种方式。Model Init Phase铺设数据结构骨架无论是通过MicroInterpreter::Invoke()的首次调用还是显式调用MicroInterpreter::AllocateTensors()在线模型分配都会开始。MicroInterpreter实例会调用MicroAllocator::StartModelAllocation()该函数开始从序列化的 flatbuffer 中拉取数据并遍历主 subgraph。根据文档及 micro_allocator.cc 的实现StartModelAllocation()按以下顺序开始分配初始化 scratch buffer 分配的内部状态清零scratch_buffer_request_count_InitScratchBufferData()并确保模型分配状态机正确model_is_allocating_置位防止重复分配根据 subgraph 中张量数量分配TfLiteEvalTensorC 结构体数组AllocateTfLiteEvalTensors()这些结构体在运行时是张量缓冲区的事实来源source of truth分配为持久分配存放在 Tail 区段为引用 flatbuffer 内联缓冲的张量赋值在这一步指向权重等内联数据的指针被直接设置为 flatbuffer 内的地址实现零拷贝复用为 subgraph 中每个算子分配NodeAndRegistrationC 结构体数组AllocateNodeAndRegistrations()每个算子对应一个NodeAndRegistration内部包含TfLiteNode与TFLMRegistration指针同样是持久分配存放在 Tail 区段若启用了压缩USE_TFLM_COMPRESSION还会在此阶段分配压缩张量列表CompressedTensorList回遍历 subgraph 算子列表用 flatbuffer 中的相关信息填充所有 C 结构体。SubgraphAllocations结构体micro_allocator.h以每个 subgraph 一个条目的方式保存node_and_registrations数组与tensorsTfLiteEvalTensor*数组是整个分配流程的核心产出。Init 阶段结束时算子内核实现已经可以接收TfLiteRegistration::init()的调用。MicroInterpreter遍历算子列表调用所有实现了该函数的算子实现。通常算子实现会返回一个对象存放到TfLiteNode结构体的user_data字段中例如各 kernel 的 OpData 结构体指针这个对象在后续推理中持续使用。Model Prepare Phase算子校验与 scratch buffer 申请在解释器完成所有算子内核的 init 之后会对 subgraph 再做一次遍历。这一次每个提供了TfLiteRegistration::prepare()函数的算子实现都会被调用。在 TFLM 中prepare 阶段用于根据模型信息校验算子的能力与参数验证张量形状shape通过TfLiteContext::GetScratchBuffer()申请任何所需的 scratch buffer计算量化运行时数据如 scale、zero point、乘子等。prepare 阶段张量的访问方式TfLiteTensor 与临时区段在 prepare 阶段算子实现会通过TfLiteTensorC 结构体请求张量数据。这个结构体比TfLiteEvalTensor更重携带更多算子初始化期间需要的信息如量化参数、维度详情等。TFLM 内部会在临时区段按请求逐个分配这些实例——临时区段正是 Arena 中 Head 与 Tail 之间的空隙。关键点prepare 阶段还没有任何数据被放置到 Head 区段因此 Head 与 Tail 之间的额外空间可用于分配缓冲区这些缓冲区的生命周期持续到MicroAllocator::ResetTempAllocations()被调用为止更多细节见 memory_management.md。重要限制TfLiteTensor结构体仅在 TFLM 的TfLiteRegistration::prepare()期间可用。此分配阶段结束后张量数据只能通过TfLiteEvalTensor结构体访问。推理热路径上使用的都是轻量的TfLiteEvalTensor这是 TFLM 性能优化的重要设计。scratch buffer 请求的登记与落位此外在 prepare 阶段每个算子实现还可以通过TfLiteContext::RequestScratchBufferInArena()请求 scratch buffer。这些请求限制为每个算子最多kMaxScratchBuffersPerOp个该常量在 micro_allocator.cc 中定义为12以internal::ScratchBufferRequest结构体记录bytes、node_idx、subgraph_idx见 micro_allocator.h的形式存储在每个算子 prepare 块对应的实例变量中实际上存放在 Head 区段开头的请求数组中当解释器切换到下一个算子时所有请求最终会被移动到 Head 区段。从源码看RequestScratchBufferInArena()micro_allocator.cc会先统计当前节点尚未赋值的请求数若超过kMaxScratchBuffersPerOp则报错返回新请求的node_idx先以kUnassignedScratchBufferRequestIndex-1作为哨兵值等到节点 prepare 结束时再真正赋上节点编号。每次TfLiteRegistration::prepare()调用完成后MicroInterpreter会调用MicroAllocator::FinishPrepareNodeAllocations()。这个方法micro_allocator.cc调用ResetTempAllocations()清理该算子 prepare 期间产生的所有临时分配扫描 scratch buffer 请求数组把哨兵值 -1 更新为真实的node_id通过ResizeBuffer()重新调整 Head 区段中 scratch buffer 请求数组的大小为下一个算子的至多kMaxScratchBuffersPerOp个请求预留空间从此开始把后续所有 scratch buffer 请求存储进 Arena 的 Head 区段。当所有算子都完成 prepare 后MicroInterpreter调用MicroAllocator::FinishModelAllocation()开始最终确定在线内存计划。Finish Model Allocation Phase提交静态内存计划在线内存规划的最后一个阶段由MicroAllocator::FinishModelAllocation()处理micro_allocator.cc。该函数执行以下任务1. 分配 scratch buffer 句柄首先通过AllocateScratchBufferHandles()在Tail区段为所有持久 scratch buffer 请求分配ScratchBufferHandle结构体数组当前位于 Head 的请求。ScratchBufferHandle只包含一个uint8_t* data指针micro_allocator.h这种指针即索引的紧凑设计保证了推理时按索引快速查找。2. 提交静态内存计划CommitStaticMemoryPlanCommitStaticMemoryPlan()micro_allocator.cc是核心中的核心实现细节如下构建 AllocationInfo使用AllocationInfoBuilder计算每个张量/缓冲区的生命周期哪些节点使用它、何时活跃并读取离线规划偏移量若模型带离线计划分配变量张量缓冲区利用此刻可用的离线规划偏移量调用AllocateVariables()在Tail区段为 variable tensor 分配持久缓冲区创建计划调用GreedyMemoryPlanner优化 Head 中的非持久空间。规划器按算子所需最大字节宽缓冲区为优化目标即按大小降序放置生成每个缓冲区在 Head 中的偏移分配指针在Tail中分配指针这些指针提供指向共享空间的指针GetOverlayMemoryAddress()以及 Head 中的偏移量设定 Head 大小根据GreedyMemoryPlanner::GetMaximumMemorySize()的结果设置 Head 区段大小ReserveNonPersistentOverlayMemory()。max_head_buffer_usage_会记录历史最大用量以支持多租户multi-tenant场景下多个模型共享 Head 区段若编译时定义了TF_LITE_SHOW_MEMORY_USE还会调用PrintMemoryPlan()输出 ASCII 布局图便于调试。3. 分配变量张量缓冲区如上所述variable tensor在多次Invoke()调用之间保留数据的张量的缓冲区最终分配在 Tail 区段。TFLM 完成在线模型分配后所有缓冲区都已就绪可进入最佳推理速度状态。此后再也不允许算子实现申请 scratch buffer——整个 Arena 布局在推理期间保持不变这也是嵌入式实时系统的确定性deterministic要求。补充如何审计 Arena 中的分配虽然online_memory_allocation_overview.md聚焦三阶段流程但结合 memory_management.md 可以顺手掌握一个实战工具TFLM 提供了记录式内存 APIRecording Memory APIs用于审计共享 Tensor Arena 的用量。只需把tflite::MicroInterpreter换成tflite::RecordingMicroInterpreter引入 recording_micro_interpreter.h再调用interpreter.GetMicroAllocator().PrintAllocations()即可输出类似如下的详细分配日志示例来自 memory_arena_threshold_test.cc[RecordingMicroAllocator] Arena allocation total 9568 bytes [RecordingMicroAllocator] Arena allocation head 7744 bytes [RecordingMicroAllocator] Arena allocation tail 1824 bytes [RecordingMicroAllocator] TfLiteEvalTensor data used 360 bytes with alignment overhead (requested 360 bytes for 15 allocations) [RecordingMicroAllocator] NodeAndRegistration struct used 392 bytes with alignment overhead (requested 392 bytes for 7 NodeAndRegistration structs) [RecordingMicroAllocator] Operator runtime data used 136 bytes with alignment overhead (requested 136 bytes for 5 OpData structs)其中各记录条目分别对应本文提到的数据结构TfLiteEvalTensor dataInit 阶段分配的评估张量结构体、NodeAndRegistration struct每个算子一个、Operator runtime dataprepare 阶段缓存的量化参数等持久数据。这对验证Head/Tail 各占多少、每个结构体花了多少非常有用。结语一次分配全程零动态内存回顾 TFLM 的在线内存分配其设计哲学可以概括为在推理开始前用一次确定性的三阶段流程Init → Prepare → Finish把 Arena 的布局彻底定死。持久结构体Tail一次成型非持久张量Head由GreedyMemoryPlanner贪心复用临时区段Temp只服务于 prepare 阶段的作用域分配flatbuffer 内联权重则零拷贝直引。推理一旦开始系统不再允许任何额外分配——这正是 TFLM 能在微控制器上提供可预测、低开销、无碎片推理的关键所在。若想继续深入建议按以下路径阅读仓库源码分配入口与状态机micro_allocator.h、micro_allocator.ccArena 底层分配器single_arena_buffer_allocator.h贪心规划算法greedy_memory_planner.h、greedy_memory_planner_test.ccArena 三段式布局总览memory_management.md离线内存计划的对照方案offline_memory_plan.md赞分享人工智能深度学习推理引擎本地部署嵌入式物联网【免费下载链接】tflite-microInfrastructure to enable deployment of ML models to low-power resource-constrained embedded targets (including microcontrollers and digital signal processors).项目地址https://gitcode.com/gh_mirrors/tf/tflite-micro点击查看免费下载相关推荐Netty 内存池申请内存全流程解析从 PooledByteBufAllocator 到 PoolChunk 的分级分配Netty 内存池申请内存全流程解析从 PooledByteBufAllocator 到 PoolChunk 的分级分配 本文基于 Netty 4.1.16文档教程技术博客知识库OceanBase 内存管理完全指南多租户内存分配接口、Arena 分配器与编程规范解析OceanBase 内存管理完全指南多租户内存分配接口、Arena 分配器与编程规范解析 导读 内存管理是 OceanBase 这类大型 C 分布式数据库数据库分布式数据库关系型数据库后端高可用arena内存区域分配管理arena内存区域分配管理 项目介绍 arena 是一个使用纯 C 语言编写的内存区域分配器实现它遵循了 stb 风格的单文件库设计模式。该项目的目标是提供上一篇SadConsole未来展望探索即将推出的令人兴奋的新功能与路线图下一篇把 Ryujinx 源码编译成能跑的模拟器手把手快速上手创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考