TensorFlowLite-Micro (TFLM)

支持的芯片

概述

TensorFlowLite-Micro(TFLM)是一个开源库,它是 TensorFlow Lite 的一个移植版本,专为数字信号处理器(DSP)、微控制器和其他内存有限的设备而设计,以便在这些设备上运行机器学习模型。

Ameba-tflite-micro 是针对 Realtek Ameba 系列芯片的 TFLM 库的优化版本,包含平台相关的优化。

代码仓库与依赖说明

代码仓库

Ameba-tflite-micro 代码托管于 GitHub:ameba-tflite-micro

依赖说明

Ameba-tflite-micro 需配合以下任一 SDK 使用:

SDK 选择取决于芯片型号与内核。

下载方式

  • 自动拉取:使用 git clone --recursive 拉取 SDK,将自动包含 ameba-tflite-micro 子模块。

  • 手动拉取:单独克隆 ameba-tflite-micro,并放置于 SDK 指定路径。

芯片

内核

SDK

TFLM 路径

RTL8721Dx/RTL8720E

KM4

ameba-rtos

{SDK}/component/tflite_micro

RTL8730E

CA32

ameba-rtos

{SDK}/component/tflite_micro

RTL8713E/RTL8726E

KM4

ameba-rtos

{SDK}/component/tflite_micro

RTL8713E/RTL8726E

DSP

ameba-dsp

{SDK}/lib/tflite_micro

示例说明

示例位于 ameba-tflite-micro 仓库的 examples 目录下。

示例

说明

tflm_hello_world

基础 TFLM 推理示例,演示浮点与量化模型的加载与推理

tflm_micro_speech

语音关键词识别 ("yes"/"no"),基于频谱图音频预处理

tflm_mnist

图像识别:手写数字分类 (0-9),使用卷积神经网络模型

编译 TFLM

编译 TFLM 示例前请确保已参考 SDK 下载 下载 XDK,并已按照 基于 CLI 的编译与下载 成功编译 SDK 默认应用。

  1. 进入 SDK 目录,运行 env.sh 设置环境变量

cd {SDK}/
source env.sh
  1. 运行 ameba.py menuconfig 进入配置界面,开启 TFLM

ameba.py menuconfig <ic_name>  ;; 请替换为芯片型号,如 RTL8721Dx

菜单路径:

----MENUCONFIG FOR General----
CONFIG TrustZone  --->
...
CONFIG APPLICATION  --->
   GUI Config  --->
   ...
   AI Config  --->
      [*] Enable TFLITE MICRO
  1. 编译示例(如需编译其他示例,将 tflm_hello_world 替换为目标示例名称)

ameba.py build -a tflm_hello_world

更多示例位于 {SDK}/component/tflite_micro/examples 目录。

示例教程

MNIST 简介

MNIST (Modified National Institute of Standards and Technology database) 是一个大型手写体数字的图片数据集,广泛用于机器学习模型的训练和验证。

本教程基于 MNIST 数据集,展示 从训练模型到部署,并在 Ameba SoC 上使用 TFLM 进行 推理 的完整流程。

示例代码: {SDK}/component/tflite_micro/examples/tflm_mnist

实验步骤

备注

步骤 1-4 是为了在开发机器(服务器或个人电脑等)上准备必要的文件。您可以跳过这些步骤,直接使用 SDK 中准备好的文件在步骤 5 中编译和烧录。

步骤 1:环境配置

进入示例目录:

cd {SDK}/component/tflite_micro/examples/tflm_mnist
  • 首次配置:创建虚拟环境并安装依赖:

python -m venv venv
source venv/bin/activate  # Windows 下使用: venv\Scripts\activate
pip install -r requirements.txt
  • 已有环境:如果已创建过虚拟环境,只需激活即可:

source venv/bin/activate  # Windows 下使用: venv\Scripts\activate

备注

由于部分库的要求,Python 版本需要在 3.8-3.11 之间。

步骤 2:训练模型

使用 Keras(TensorFlow)或 PyTorch 训练 MNIST 10 个数字的分类模型。

示例所用模型是一个以卷积为主的简单网络结构,在数据集上进行几轮训练迭代后,测试它的准确率。

运行脚本

python keras_train_eval.py --output keras_mnist_conv

输出文件:

  • keras_mnist_conv/saved_model/ - Keras SavedModel 格式

备注

由于微处理器的 算力内存 有限,建议关注 模型大小 以及 运算操作数。在 keras_train_eval.py 脚本中,使用了 keras_flops 库:

from keras_flops import get_flops

model.summary()
flops = get_flops(model, batch_size=1)

步骤 3:转换成 Tflite 文件

对训练好的模型应用 训练后整数量化,输出为 .tflite 格式文件。

Ameba SoC 也支持浮点模型的推理,但仍建议使用整数量化,因为它可以在性能下降很小的情况下,极大地减少计算量和内存。

有关全整数量化的更多详细信息,请参考 tflite 官网

运行脚本

python convert.py --input-path keras_mnist_conv/saved_model --output-path keras_mnist_conv

备注

convert.py 脚本细节说明:

  • 使用 tf.lite.TFLiteConverter 转换 SavedModel 为 .tflite 格式

  • 通过以下设置实现 int8 量化:

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = repr_dataset
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8
tflite_int8_model = converter.convert()

转换完成后,示例将使用 TFLite 模型验证测试集性能。

输出文件:

  • {output}/model_int8.tflite - Int8 量化后的 TFLite 模型

  • testdata/input_int8.npy - 测试输入数据(100 个样本)

  • testdata/label_int8.npy - 测试标签数据(100 个样本)

步骤 4:转换成 C++ 并准备部署

运行以下命令,把 .tflite 模型和 .npy 测试数据转换成部署所需的 .cc 和 .h 格式文件:

cp {output}/model_int8.tflite models/model_int8.tflite
python generate_cc_arrays.py models models/model_int8.tflite
python generate_cc_arrays.py testdata testdata/input_int8.npy testdata/input_int8.npy testdata/label_int8.npy testdata/label_int8.npy

输出文件:

  • models/model_int8_model_data.{cc,h} - 模型数据

  • testdata/input_int8_test_data.{cc,h} - 测试输入数据

  • testdata/label_int8_test_data.{cc,h} - 测试标签数据

备注

  1. 如需使用自己训练的模型,在编译前将生成的文件替换 models/testdata/ 目录下的对应文件即可。

  2. 生成的文件将用于 example_tflm_mnist.cc 在芯片上运行推理,计算准确率,以及统计内存和延迟情况。

  3. 可使用 netron 来可视化 .tflite 文件并查看模型使用的算子。实现推理时,需实例化操作解析器来注册所需算子:

using MnistOpResolver = tflite::MicroMutableOpResolver<4>;

TfLiteStatus RegisterOps(MnistOpResolver& op_resolver) {
    TF_LITE_ENSURE_STATUS(op_resolver.AddFullyConnected());
    TF_LITE_ENSURE_STATUS(op_resolver.AddConv2D());
    TF_LITE_ENSURE_STATUS(op_resolver.AddMaxPool2D());
    TF_LITE_ENSURE_STATUS(op_resolver.AddReshape());
    return kTfLiteOk;
}
  1. 有关使用 TFLM 进行推理的更多详细信息,请参考 tflite-micro 官方网站

步骤 5:编译和烧录固件

按照 编译 TFLM 中的步骤编译固件,使用 Flash 烧写工具 烧录固件。

步骤 6:查看日志

烧录完成后,打开串口终端查看推理结果:

[TFLITE-MICRO] ~~~TESTS START~~~
[TFLITE-MICRO] model arena size = 16204
[TFLITE-MICRO] Accuracy: 100/100

"Unique Tag","Total ticks across all events with that tag."
CONV_2D, 29763
MAX_POOL_2D, 3422
RESHAPE, 339
FULLY_CONNECTED, 28897
"total number of ticks", 62421

[RecordingMicroAllocator] Arena allocation total 16204 bytes
[RecordingMicroAllocator] Arena allocation head 13968 bytes
[RecordingMicroAllocator] Arena allocation tail 2236 bytes

[TFLITE-MICRO] Total Time: 72.185 ms

[TFLITE-MICRO] ~~~ALL TESTS PASSED~~~

日志说明:

  • model arena size = 16204 - 模型所需的内存池大小(字节)

  • Accuracy: 100/100 - 100 个测试样本中正确预测的数量

  • CONV_2DMAX_POOL_2D 等 - 各算子类型消耗的时间(微秒)

  • total number of ticks = 62421 - 总推理 tick 数

  • Total Time: 72.185 ms - 总推理时间(毫秒)

备注

推理时间和内存占用可能因芯片型号(如 RTL8721Dx、RTL8730E)、内核(KM4、CA32、DSP)和 SDK 版本的不同而有所差异。以上数值仅供参考。