当前位置：首页 > news >正文

在 48KB SRAM 里挤出每一字节：TFLite Micro MemoryPlanner 的 Greedy 与 Linear 分配策略深度解读与取舍分析

news 2026/7/8 15:04:47

在 48KB SRAM 里挤出每一字节：TFLite Micro MemoryPlanner 的 Greedy 与 Linear 分配策略深度解读与取舍分析

一、MCU 推理的终极约束：不是算力不够，而是内存不足

在嵌入式端部署神经网络时，内存是最稀缺的资源。一块典型的 Cortex-M4 芯片拥有 168MHz 主频和 128KB SRAM，看似足以运行 MobileNetV1-0.25（约 200KB Flash + 80KB SRAM）。但实际分配中，SRAM 需要同时容纳激活张量、权重缓存、DMA 缓冲区和 RTOS 栈——80KB 中的 60KB 必须分配给激活张量，留给其他组件仅 20KB。

TFLite Micro（TFLM）通过"内存竞技场"（Memory Arena）机制来解决这一问题。它将所有运行时内存需求整合到一个预分配的连续缓冲区中，由 MemoryPlanner 负责在推理的不同阶段为每个张量分配和回收内存。这种方式避免了动态内存分配的不确定性和碎片化，将 SRAM 使用量从"所有张量同时驻留"降低到"同一时刻仅活跃张量驻留"。

对于一个 20 层的小型图像分类网络，线性分配可能需要 120KB，而 Greedy 内存复用可以将峰值降至 48KB——60% 的节省。这个差距在 128KB SRAM 的平台上决定了模型能否部署。

二、MemoryPlanner 的核心算法：从离线分配到在线复用

TFLM 使用离线内存规划策略。在模型解释器初始化阶段，所有张量的生命周期（从哪个算子产生、在哪个算子被最后一次消费）已通过图拓扑排序完全确定。MemoryPlanner 的任务是为这些生命周期已知的张量分配 Arena 中的偏移，使得峰值占用最小化。

flowchart TD A["解析 FlatBuffer 模型图"] --> B["执行图拓扑排序<br/>确定算子执行顺序"] B --> C["为每个张量计算<br/>first_used_op 和 last_used_op<br/>（生命周期跨度）"] C --> D{"选择分配策略"} D -- "Greedy" --> E["GreedyMemoryPlanner"] D -- "Linear" --> F["LinearMemoryPlanner"] subgraph Greedy["Greedy 策略"] E --> E1["按张量 size 降序排序"] E1 --> E2["遍历排序列表，为每个张量<br/>寻找最早的可复用偏移"] E2 --> E3["检查候选偏移处<br/>是否有存活冲突（生命周期重叠）"] E3 -- 无冲突 --> E4["分配该偏移"] E3 -- 有冲突 --> E5["增加偏移，继续搜索"] E4 --> E6["更新 peak_arena_size"] E5 --> E3 end subgraph Linear["Linear 策略"] F --> F1["维护一个可用缓冲区偏移列表"] F1 --> F2["按算子顺序遍历张量"] F2 --> F3["分配时：若偏移列表中有足够<br/>空间且前一个占用者已失效，<br/>直接复用"] F3 --> F4["分配时：若无可用空间，<br/>在当前最高偏移后追加"] F2 --> F5["释放时：将释放的偏移<br/>放回可用列表"] end E6 --> G["输出 Arena 大小"] F4 --> G G --> H["解释器根据分配表<br/>为每个张量设置 data 指针"]

2.1 Greedy 策略的离线性质与冲突检测

Greedy 策略基于区间图着色问题的离线解法。它将张量按大小降序排序，每次尝试将最大的张量分配到 Arena 的最低可用位置。冲突检测通过检查新分配张量的生命周期是否与当前位置的所有已分配张量重叠来实现。

这个策略之所以有效，是因为神经网络中通常是少数大张量（特征图）占据大部分内存，而大量小张量（一维偏置、scale 因子）可以灵活地填入大张量之间的空隙。先分配大张量、再填充小张量的策略能最大限度地减少碎片。

2.2 张量生命周期的精确计算

生命周期的计算依赖对图结构的精确分析。一个张量在以下情况变为"活跃"：

作为常量输入：在模型加载时就存活
作为算子的输出：在该算子执行后变为活跃

一个张量在以下情况变为"不活跃"（可被回收）：

所有消费它的算子都已执行完毕

这个分析看似简单，但在存在分支（如残差连接）的图中复杂度显著增加。跳跃连接中的张量需要保持活跃直到 Add 算子执行——这意味着一组张量需要跨越 10-20 个算子的执行周期，在此期间它们持续占用 Arena 空间。

三、生产级代码实现：简化版 MemoryPlanner 原型

以下展示一个简化的 Greedy MemoryPlanner 实现，揭示其核心逻辑：

#include <stdint.h> #include <stdbool.h> #include <stdlib.h> #include <string.h> #include <stdio.h> /* ================================================================ * 数据结构定义 * ================================================================ */ /* 张量信息 */ typedef struct { int32_t tensor_id; /* 张量唯一标识 */ int32_t first_used_op; /* 首次被使用的算子索引 */ int32_t last_used_op; /* 最后被使用的算子索引 */ size_t size_bytes; /* 张量大小（字节） */ bool is_constant; /* 是否为常量张量（不可复用） */ } tensor_lifetime_t; /* Arena 中的分配记录 */ typedef struct { int32_t tensor_id; int32_t first_op; int32_t last_op; size_t offset; /* 在 Arena 中的偏移 */ size_t size_bytes; } arena_allocation_t; /* MemoryPlanner 上下文 */ typedef struct { tensor_lifetime_t *tensors; /* 张量生命周期数组 */ int32_t num_tensors; /* 张量总数 */ arena_allocation_t *allocations; /* 分配结果 */ int32_t num_allocations; size_t arena_size; /* 最终需要的 Arena 大小 */ } planner_context_t; /* ================================================================ * 辅助函数：检查两个张量的生命周期是否重叠 * 重叠判定：不是 (A 在 B 开始前结束) 且不是 (B 在 A 开始前结束) * ================================================================ */ static bool lifetimes_overlap( int32_t a_first, int32_t a_last, int32_t b_first, int32_t b_last) { return !(a_last < b_first || b_last < a_first); } /* ================================================================ * 辅助函数：检查在指定偏移处是否可以分配张量 * 将新分配的 lifetime 与所有已分配的 lifetime 做冲突检测。 * ================================================================ */ static bool can_allocate_at( planner_context_t *ctx, const tensor_lifetime_t *tensor, size_t proposed_offset) { for (int32_t i = 0; i < ctx->num_allocations; i++) { arena_allocation_t *existing = &ctx->allocations[i]; /* 检查生命周期是否重叠 */ if (lifetimes_overlap( tensor->first_used_op, tensor->last_used_op, existing->first_op, existing->last_op)) { /* 生命周期重叠 → 检查偏移是否冲突 */ size_t exist_end = existing->offset + existing->size_bytes; size_t prop_end = proposed_offset + tensor->size_bytes; /* 任意重叠判定：不是 (A 在 B 结束之后) 也不是 (B 在 A 结束之后) */ if (!(proposed_offset >= exist_end || exist_end <= proposed_offset)) { /* 实际检查：两区间有交集 */ if (!(proposed_offset >= exist_end || existing->offset >= prop_end)) { return false; /* 空间冲突，无法在此处分配 */ } } } } return true; /* 没有冲突 */ } /* ================================================================ * Greedy Memory Planner 主算法 * * 算法步骤： * 1. 按张量大小降序排序（先分配大张量） * 2. 对每个张量，从 offset=0 开始搜索第一个无冲突位置 * 3. 计算最终 Arena 所需的总大小 * * 时间复杂度：O(n² × k)，n=张量数，k=平均搜索步数 * 对于 n≤200 的模型（TinyML 场景），性能可接受。 * ================================================================ */ /* qsort 比较函数：按 size 降序排列 */ static int cmp_size_desc(const void *a, const void *b) { const tensor_lifetime_t *ta = (const tensor_lifetime_t *)a; const tensor_lifetime_t *tb = (const tensor_lifetime_t *)b; if (ta->size_bytes > tb->size_bytes) return -1; if (ta->size_bytes < tb->size_bytes) return 1; return 0; } int greedy_memory_planner(planner_context_t *ctx) { if (!ctx || !ctx->tensors || !ctx->allocations || ctx->num_tensors == 0) { fprintf(stderr, "[Planner] 无效的上下文参数\n"); return -1; } /* 复制张量列表并排序（不修改原始数据） */ tensor_lifetime_t *sorted = (tensor_lifetime_t *)malloc( ctx->num_tensors * sizeof(tensor_lifetime_t)); if (!sorted) { fprintf(stderr, "[Planner] 内存分配失败\n"); return -2; } memcpy(sorted, ctx->tensors, ctx->num_tensors * sizeof(tensor_lifetime_t)); qsort(sorted, ctx->num_tensors, sizeof(tensor_lifetime_t), cmp_size_desc); ctx->num_allocations = 0; ctx->arena_size = 0; /* 逐张量寻找最佳位置 */ for (int32_t i = 0; i < ctx->num_tensors; i++) { tensor_lifetime_t *tensor = &sorted[i]; bool allocated = false; /* 从 offset=0 开始尝试 */ for (size_t offset = 0; !allocated && offset <= ctx->arena_size; offset += 4) { /* 步进 4 字节确保对齐 */ if (can_allocate_at(ctx, tensor, offset)) { /* 找到合适位置 —— 记录分配 */ ctx->allocations[ctx->num_allocations].tensor_id = tensor->tensor_id; ctx->allocations[ctx->num_allocations].first_op = tensor->first_used_op; ctx->allocations[ctx->num_allocations].last_op = tensor->last_used_op; ctx->allocations[ctx->num_allocations].offset = offset; ctx->allocations[ctx->num_allocations].size_bytes = tensor->size_bytes; ctx->num_allocations++; /* 更新 Arena 峰值大小 */ size_t end_offset = offset + tensor->size_bytes; if (end_offset > ctx->arena_size) { ctx->arena_size = end_offset; } allocated = true; } } if (!allocated) { /* 所有位置都有冲突 → 在 Arena 末尾追加 */ size_t end_offset = ctx->arena_size; ctx->allocations[ctx->num_allocations].tensor_id = tensor->tensor_id; ctx->allocations[ctx->num_allocations].first_op = tensor->first_used_op; ctx->allocations[ctx->num_allocations].last_op = tensor->last_used_op; ctx->allocations[ctx->num_allocations].offset = end_offset; ctx->allocations[ctx->num_allocations].size_bytes = tensor->size_bytes; ctx->num_allocations++; ctx->arena_size = end_offset + tensor->size_bytes; } } free(sorted); /* 验证：Arena 大小应为 16 字节对齐（ARM CMSIS-NN 的要求） */ ctx->arena_size = (ctx->arena_size + 15) & ~15UL; return 0; } /* ================================================================ * 使用示例和打印函数 * ================================================================ */ void print_allocation_plan(const planner_context_t *ctx) { if (!ctx) return; printf("=== 内存分配计划 ===\n"); printf("Arena 总大小: %zu 字节\n", ctx->arena_size); printf("%-12s %-8s %-8s %-8s %-10s\n", "TensorID", "FirstOp", "LastOp", "Offset", "Size(B)"); printf("----------------------------------------\n"); for (int32_t i = 0; i < ctx->num_allocations; i++) { arena_allocation_t *a = &ctx->allocations[i]; printf("%-12d %-8d %-8d %-8zu %-10zu\n", a->tensor_id, a->first_op, a->last_op, a->offset, a->size_bytes); } }

四、边界分析与架构权衡：Greedy 与 Linear 的策略对比

Greedy 策略的核心假设是"按大小排序能找到接近最优的解"。在大多数神经网络中这个假设成立，但在以下边界情况下会退化：

残差网络中的长生命周期张量：ResNet 的跳跃连接要求输入张量存活到 Add 算子（跨越 10+ 步）。这意味着一个 64KB 的张量可能需要在 Arena 中保留 10 个算子的执行时间。Greedy 策略会尝试找到不与这个"长期租户"重叠的位置，但在 Arena 空间紧张时可能导致碎片化加剧——大张量占据低偏移区，迫使后续中等大小的张量堆叠在高偏移区。

作为对比，Linear 策略按算子执行顺序分配内存。它更像一个"先来先服务"的系统——按图拓扑顺序逐个处理算子，为每个输入/输出张量分配空间。这种策略在长生命周期张量较少的网络中效果最好。

两种策略在 MobileNetV1-0.25（20 层，64KB SRAM 目标）上的对比：