TensorFlow Lite Micro: Embedded Machine Learning on TinyML Systems

该论文主要介绍如何将已经训练完成的神经网络模型转换并部署到微控制器等资源受限设备中，以及 TFLM 为解决嵌入式设备碎片化、内存不足和算子性能差异所采用的设计方法。

1. 对 TFLM 定位的理解

模型训练仍然在计算资源较充足的计算机或服务器上完成。TFLM 的主要作用是接收经过 TensorFlow Lite 转换的模型，并在 MCU 上完成模型加载、内存分配、算子调用和推理执行。

因此，该论文所描述的 TinyML 流程可以概括为：

计算机端采集和处理数据
        ↓
使用 TensorFlow / Keras 训练模型
        ↓
TensorFlow Lite 转换和优化
        ↓
生成 .tflite FlatBuffer 模型
        ↓
将模型写入 MCU 程序
        ↓
TFLM 初始化模型、算子和内存
        ↓
输入传感器数据
        ↓
执行模型推理
        ↓
输出分类或检测结果
        ↓
测量内存、延迟和运行周期

2. TFLM 的 TinyML 实现流程

2.1 在高性能平台上训练模型

模型首先在 TensorFlow 训练环境中完成设计和训练。训练阶段通常在计算机、服务器或 GPU 上进行，而不是直接在 MCU 上完成。

这一阶段产生的是 TensorFlow 训练图，其中包含：

• 网络结构；
• 模型权重；
• 输入和输出；
• 训练阶段使用的算子；
• Batch Normalization、Dropout 等训练组件。

需要注意的是，TensorFlow 支持的算子数量远多于 TensorFlow Lite 和 TFLM。因此，能够在计算机上训练的模型不一定能够直接部署到 MCU。论文指出，开发者有时会在完成模型训练后才发现模型包含目标推理框架不支持的算子。

2.2 使用 TensorFlow Lite 转换模型

模型训练完成后，通过 TensorFlow Lite Converter 将训练转换为推理。

论文图 1 给出的模型导出流程为：

TensorFlow 训练
        ↓
TensorFlow Lite 导出器
        ↓
推理
        ↓
有序算子列表 + 模型权重
        ↓
TensorFlow Lite FlatBuffer 文件

转换阶段可以进行以下处理：

• 将浮点模型转换为 8 bit 等量化表示；
• 将常量表达式提前计算并固化；
• 将 Batch Normalization 合并到相邻算子；
• 删除只在训练期间使用的 Dropout；
• 将模型转换为 TensorFlow Lite 支持的算子集合。

TFLM 直接复用 TensorFlow Lite 的模型转换和优化工具链，转换结果是一个 FlatBuffer 格式的 .tflite 文件。

2.3 将模型存储到嵌入式程序中

许多 MCU 没有文件系统，不能像计算机一样直接读取 .tflite 文件。因此，模型通常需要转换为 C 或 C++ 数组，再与应用程序一起编译到固件中。

TFLM 使用 FlatBuffer 作为模型序列化格式。该格式可以直接从内存中读取，不需要先解包为另一种数据结构。对于没有文件系统的设备，FlatBuffer 文件可以转换成包含模型数据的 C 源文件，再编译到最终程序中。

2.4 创建模型对象和算子解析器

在 MCU 程序中，TFLM 首先需要创建模型对象。

随后创建 OpResolver，即算子解析器。其作用是声明模型运行时需要使用哪些算子，例如：

• 卷积；
• 全连接；
• 池化；
• 激活函数；
• Softmax；
• 量化相关算子。

OpResolver 只将模型真正需要的算子链接到最终二进制程序中，从而减少程序文件大小。

2.5 创建固定内存区域 Tensor Arena

MCU 通常缺少完整的动态内存管理机制。为避免运行过程中频繁申请和释放内存，TFLM 要求应用程序预先提供一块连续的固定内存，称为 Tensor Arena。

Tensor Arena 用于存储：

• 输入和输出张量；
• 中间计算结果；
• 算子状态；
• 模型元数据；
• 临时计算缓冲区；
• 持久化变量。

TFLM 在初始化阶段从 Tensor Arena 中分配所有需要的内存，初始化完成后不再进行动态分配，从而避免堆碎片对长时间运行系统造成影响。

2.6 执行内存规划

不同中间张量的生命周期并不相同。例如，一个卷积层输出只需要保留到后续使用它的算子执行完成，之后这部分内存就可以被其他张量复用。

TFLM 的 Memory Planner 会分析：

• 每个中间张量的大小；
• 张量从何时开始使用；
• 张量在何时不再需要；
• 哪些张量的内存可以重叠复用。

论文采用类似 Bin Packing 的方法压缩中间张量所需的内存空间。其基本过程是按照临时内存块大小排序，并将其放入能够容纳它的已有空闲区域中，从而减少 Tensor Arena 的总大小。

2.7 创建 TFLM 解释器

完成模型、算子解析器和 Tensor Arena 的准备后，系统创建 TFLM 解释器。

解释器初始化时接收：

模型对象
+ OpResolver
+ Tensor Arena

解释器负责：

• 加载模型结构；
• 将模型算子映射到具体实现；
• 计算张量内存位置；
• 初始化算子；
• 按照模型中的算子顺序执行计算。

论文选择解释器而不是完全生成固定机器代码，主要原因是解释器具有更好的模型可移植性和维护性。模型结构和权重与执行代码相对分离，更新模型时不需要完全重写整个推理框架。

2.8 输入传感器数据并执行推理

模型初始化完成后，应用程序获得输入张量的内存地址，并将传感器数据写入输入张量。

写入输入后，应用程序调用解释器执行推理。解释器按照拓扑排序后的算子列表逐个执行计算，并通过 Memory Planner 预先计算的内存偏移找到每个算子的输入和输出。

模型运行结束后，应用程序读取输出张量，获得分类概率、预测类别或异常分数。

2.9 使用硬件专用算子进行加速

TFLM 默认提供可移植的参考算子实现，但参考实现不一定能够充分利用特定 MCU 的指令集。

因此，TFLM 允许硬件厂商使用优化后的算子替换默认实现。例如，Arm 平台可以使用 CMSIS-NN 提供的优化算子，包括：

• 卷积；
• 激活；
• 全连接；
• 池化；
• Softmax；
• 基础数学运算。

TFLM 通过统一的算子 API 隐藏不同硬件实现的细节。模型和解释器结构保持不变，只替换计算密集型算子的底层实现，从而兼顾可移植性与性能。

2.10 对系统进行测试与评价

论文并没有只使用模型准确率评价 TinyML 系统，还测量了：

• 模型运行周期；
• 算子计算周期；
• 解释器运行开销；
• 持久内存；
• 临时内存；
• 总内存占用；
• 不同硬件上的性能；
• 优化算子与参考算子的性能差异。

论文实验表明，TFLM 解释器本身的代码占用低于 2 KB。不同模型的内存需求差异明显，表中卷积参考模型、Google Hotword 和 VWW 模型的总内存分别约为 9.04 KB、12.80 KB 和 81.79 KB。

TFLM 还提供 Benchmark 和 Profiling API，用于比较硬件平台、测量执行性能并定位耗时较高的算子。其基准测试后来也被 tinyMLPerf 使用。

3. 本次阅读形成的关键认识

通过这篇论文，我对 TinyML 的流程形成了以下认识。

3.1 TinyML 不等于直接训练一个小模型

TinyML 实际上是一条完整的工程链路：

模型设计
→ 训练
→ 量化和转换
→ 算子兼容
→ 内存规划
→ 嵌入式运行时
→ 硬件优化
→ 性能测试

模型参数量较小只是其中一个条件。即使模型本身较小，如果包含 TFLM 不支持的算子、需要大量中间张量或依赖动态输入，也可能无法部署。

3.2 模型结构需要考虑部署框架

模型设计和部署不能完全分开。在训练模型时需要提前检查：

• 算子是否受到 TFLM 支持；
• 是否能够进行 INT8 量化；
• 是否包含动态形状；
• 中间特征图是否占用过多 RAM；

3.3 MCU 的主要限制不仅是模型文件大小

TinyML 系统同时受到两类内存限制：

• Flash/ROM：存放模型权重、算子代码和程序；
• SRAM/RAM：存放输入、输出、中间张量和临时缓冲区。

部分模型的权重可能较小，但中间特征图较大，因此峰值 RAM 仍可能成为部署瓶颈。

3.4 量化主要发生在模型转换阶段

TFLM 本身主要负责运行量化后的模型。FP32 到 INT8 的转换通常由 TensorFlow Lite Converter 完成，部分量化方法还需要在训练阶段加入量化感知训练。