3. scratch buffer 设计原理：为什么需要 scratch buffer、生命周期、与 tensor arena 的关系
TFLite Micro 的推理引擎在嵌入式设备上运行时，面临的最核心约束不是算力，而是内存。一个典型的 Cortex-M4 设备可能只有 192KB RAM，而一个 MobileNetV1 模型仅权重就需要约 4.2MB。显然，模型权重和中间张量不可能全部常驻内存。TFLite Micro 通过 tensor arena（张量竞技场）统一管理中间张量，而 scratch buffer（临时缓冲区）则是 arena 内部一种特殊的、由算子主动申请和释放的临时内存块。理解 scratch buffer 的设计原理，是排查“申请失败”类问题的前提——只有知道它为什么存在、何时存活、与 arena 如何交互，才能准确判断失败发生在哪一层。
3.1 为什么需要 scratch buffer：从算子执行的内存需求说起
在 TFLite Micro 的推理流程中，每个算子（Operator）的 Invoke() 函数负责完成一次张量到张量的计算。但很多算子并非“输入直接写到输出”那么简单。以卷积为例，为了提升计算效率，通常需要将输入数据按卷积核窗口重新排列成矩阵（im2col），这个中间矩阵的大小可能远超输入张量本身。又比如某些量化算子需要临时的浮点累加器，或者某些融合算子需要暂存中间结果。
这些临时内存有两个显著特点：大小在编译期无法确定（取决于输入尺寸、算子参数），生命周期极短（仅在当前算子执行期间有效）。如果为每个算子都预留一块独立内存，arena 会迅速膨胀；如果直接复用输出张量内存，又可能破坏尚未读取的输入数据。scratch buffer 正是为解决这一矛盾而设计的——它允许算子在执行前向 MicroAllocator 申请一块临时内存，执行完毕后立即释放，供下一个算子复用。
3.2 scratch buffer 的生命周期：从 Request 到 Release
scratch buffer 的生命周期严格绑定在单个算子的执行周期内，具体分为四个阶段：
- 声明阶段（Prepare）：在算子的
Prepare()函数中，通过context->RequestScratchBufferInArena(&context, bytes, &buffer_index)请求一块大小为bytes的缓冲区。此时只是“登记”需求，并未真正分配内存。 - 分配阶段（Allocate）：所有算子的
Prepare()执行完毕后，TFLite Micro 的MicroAllocator统一计算所有 scratch 请求的总大小，在 arena 的临时区尾部一次性分配。 - 使用阶段（Invoke）：算子执行时，通过
context->GetScratchBuffer(&context, buffer_index)获取实际内存指针，进行读写操作。 - 释放阶段（Reset）：算子执行完毕后，scratch buffer 并不会立即释放，而是由
MicroAllocator在下一轮推理开始前统一重置。这意味着同一块 scratch 内存可以被不同算子在不同推理轮次中复用。
| 阶段 | 调用时机 | 关键 API | 内存状态 |
|---|---|---|---|
| 声明 | Prepare() | RequestScratchBufferInArena | 仅记录大小，不分配 |
| 分配 | 所有 Prepare 完成后 | MicroAllocator 内部处理 | 在临时区尾部预留空间 |
| 使用 | Invoke() | GetScratchBuffer | 可读写 |
| 重置 | 推理轮次开始前 | ResetTempAllocations | 标记为可复用 |
RequestScratchBufferInArena 只是登记需求，真正的内存分配发生在所有 Prepare() 完成后。如果在 Prepare() 阶段之后才调用 RequestScratchBufferInArena，会导致分配失败或内存重叠。
3.3 scratch buffer 与 tensor arena 的关系：子区域与复用机制
从内存布局上看，tensor arena 是一个连续的大数组，由 MicroAllocator 统一管理。arena 内部划分为两个逻辑区域：静态区（存放权重、常量张量、模型输入输出）和临时区（存放中间激活值）。scratch buffer 位于临时区的尾部，与临时张量共享同一块内存池。
这种设计带来一个关键特性：scratch buffer 与临时张量可以相互覆盖。例如，当算子 A 执行完毕，其输出张量（临时区中的一块内存）不再需要时，算子 B 的 scratch buffer 可以复用这块内存。TFLite Micro 通过 MicroAllocator 的分配记录来跟踪哪些内存块是“活跃”的，哪些是“可回收”的。
理解这一关系对排查“申请失败”至关重要：
- arena 总大小固定：由
MicroMutableOpResolver和tensor_arena_size参数决定，通常在模型初始化时传入。 - 静态区优先分配：模型加载时，所有持久化张量先占用 arena 头部空间。
- 临时区动态分配：推理过程中，临时张量和 scratch buffer 在剩余空间中“挤占”内存。
- 失败原因：当剩余空间不足以满足某个算子的 scratch 请求时，
MicroAllocator返回kTfLiteError，导致推理中断。
RecordingMicroAllocator 获取实际峰值内存，再调整 arena 大小。
3.4 实战场景：一个典型的 scratch 申请失败案例
假设我们在一款基于 STM32F746（320KB RAM）的设备上部署一个语音唤醒模型，arena 设置为 200KB。模型包含 5 个深度可分离卷积层和一个全连接层。初始现象是：模型加载成功，前 3 层推理正常，执行到第 4 层时返回 kTfLiteError，日志显示 Failed to allocate scratch buffer。
分析过程：
- 使用
RecordingMicroAllocator记录内存分配，发现第 4 层卷积的 im2col 缓冲区需要 48KB，而当前 arena 剩余空间仅 32KB。 - 进一步检查发现，第 3 层的输出张量（16KB）虽然在第 4 层执行前已不再需要，但 TFLite Micro 的临时张量复用策略是“按需分配”，不会主动回收已释放的临时张量。
- 真正的问题在于：第 4 层的输入张量（第 3 层输出）与 scratch buffer 同时存在，而 arena 的临时区被前几层的中间结果碎片化。
解决方案：
- 将 arena 从 200KB 增加到 240KB，为 scratch 预留足够空间。
- 优化模型结构：将第 4 层的
dilation_rate从 2 改为 1，减少 im2col 缓冲区大小（从 48KB 降至 24KB）。 - 在代码中显式调用
context->RequestScratchBufferInArena时，将大缓冲区拆分为多个小缓冲区，提高内存碎片利用率。
结果对比：调整后，峰值内存从 198KB 降至 176KB，推理成功，且推理速度提升 12%（因为 im2col 数据量减少）。
3.5 常见误区与设计建议
在实际项目中，开发者对 scratch buffer 的理解常存在以下误区：
| 误区 | 实际情况 | 正确做法 |
|---|---|---|
| 认为 scratch 越大越好 | 过大的 scratch 会挤压临时张量空间，导致其他算子失败 | 按需申请，尽量复用已有缓冲区 |
| 在 Invoke 中动态申请 | Invoke 阶段不允许申请新 scratch | 所有申请必须在 Prepare 阶段完成 |
| 忽略对齐要求 | 某些硬件（如 SIMD）要求 16 字节对齐 | 使用 AlignTo 宏确保对齐 |
| 认为 scratch 是全局共享的 | 每个算子只能访问自己申请的 buffer_index | 通过 context 传递 buffer_index，不可跨算子访问 |
设计建议：在编写自定义算子时，优先考虑是否可以使用已有的临时张量替代 scratch buffer。例如，如果算子的中间结果可以写入输出张量（在确认不会破坏输入的前提下），就无需额外申请 scratch。此外，对于多个算子共用的临时缓冲区，可以在模型层面合并为一个大的 scratch 请求，减少分配次数和碎片。