跳到主要内容
版本:0.10.0

支持 TFLite 模型的 microTVM

备注

单击 此处 下载完整的示例代码

作者:Tom Gall

本教程介绍如何用 microTVM 和支持 Relay 的 TFLite 模型。

备注

若要在 microTVM 虚拟机参考手册上运行本教程,请根据页面底部的链接下载 Jupyter Notebook 并将其保存到 TVM 目录中。然后:

  1. 用修改后的 vagrant ssh 命令登录到虚拟机参考手册: $ vagrant ssh -- -L8888:localhost:8888
  2. 安装 Jupyter: pip install jupyterlab
  3. cd 到 TVM 目录
  4. 安装 TFLite:poetry install -E importer-tflite
  5. 启动 Jupyter Notebook:jupyter notebook
  6. 复制 localhost URL,并将其粘贴到浏览器中
  7. 导航到已保存的 Jupyter Notebook(.ipynb 文件)

设置​

安装 TFLite​

开始前,先安装 TFLite 包,两种安装方式如下所示:

  1. 使用 pip 安装

    pip install tflite=2.1.0 --user
  2. 生成 TFLite 包,步骤如下:

    获取 flatc 编译器,参阅 https://github.com/google/flatbuffers 了解详细信息,确保已正确安装。

    flatc --version

    获取 TFLite schema。

    wget https://raw.githubusercontent.com/tensorflow/tensorflow/r1.13/tensorflow/lite/schema/schema.fbs

    生成 TFLite 包。

    flatc --python schema.fbs

    将当前文件夹(包含生成的 TFLite 模块)添加到 PYTHONPATH。

    export PYTHONPATH=${PYTHONPATH:+$PYTHONPATH:}$(pwd)

用 python -c "import tflite" 验证 TFLite 包是否安装成功。

安装 Zephyr(仅限物理硬件)​

用主机模拟运行本教程时(默认),使用主机 gcc 构建模拟设备的固件镜像。若编译到物理硬件上运行,需安装一个 toolchain 以及一些特定于 target 的依赖。microTVM 允许用户提供任何可以启动 TVM RPC 服务器的编译器和 runtime。开始之前,请注意本教程依赖于 Zephyr RTOS 来提供这些部分。

参考 安装说明 安装 Zephyr。

题外话:重新创建预训练 TFLite 模型

本教程下载了预训练的 TFLite 模型。使用微控制器时,请注意这些设备的资源高度受限,像 MobileNet 这样的标准模型和小内存并不匹配。

本教程使用 TF Micro 示例模型之一。

若要复制训练步骤,参阅:https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/micro/examples/hello_world/train

备注

若不小心从 wget https://storage.googleapis.com/download.tensorflow.org/models/tflite/micro/hello_world_2020_04_13.zip 下载了示例预训练模型,会由于未实现的操作码(114)而失败。

加载并准备预训练模型​

把预训练 TFLite 模型从当前目录中的文件中加载到 buffer

import os
import json
import tarfile
import pathlib
import tempfile
import numpy as np

import tvm
from tvm import relay
import tvm.contrib.utils
from tvm.contrib.download import download_testdata

use_physical_hw = bool(os.getenv("TVM_MICRO_USE_HW"))
model_url = "https://people.linaro.org/~tom.gall/sine_model.tflite"
model_file = "sine_model.tflite"
model_path = download_testdata(model_url, model_file, module="data")

tflite_model_buf = open(model_path, "rb").read()

利用 buffer,转换为 TFLite 模型 Python 对象

try:
import tflite

tflite_model = tflite.Model.GetRootAsModel(tflite_model_buf, 0)
except AttributeError:
import tflite.Model

tflite_model = tflite.Model.Model.GetRootAsModel(tflite_model_buf, 0)

打印模型的版本

version = tflite_model.Version()
print("Model Version: " + str(version))

输出结果:

Model Version: 3

解析 Python 模型对象,并转换为 Relay 模块和权重。注意输入张量的名称必须与模型中包含的内容相匹配。

若不确定,可通过 TensorFlow 项目中的 visualize.py 脚本来查看。参阅 如何检查 .tflite 文件?

input_tensor = "dense_4_input"
input_shape = (1,)
input_dtype = "float32"

mod, params = relay.frontend.from_tflite(
tflite_model, shape_dict={input_tensor: input_shape}, dtype_dict={input_tensor: input_dtype}
)

定义 target​

接下来为 Relay 创建一个构建配置,关闭两个选项,然后调用 relay.build,为选定的 TARGET 生成一个 C 源文件。

当在与主机( Python 脚本执行的位置)相同架构的模拟 target 上运行时,为 TARGET 选择下面的「host」,选择 C Runtime 作为 RUNTIME ,并选择适当的单板/虚拟机来运行它(Zephyr 将创建基于 BOARD 的正确 QEMU 虚拟机)。

下面的示例中,选择 x86 架构并相应地选择 x86 虚拟机:

RUNTIME = tvm.relay.backend.Runtime("crt", {"system-lib": True})
TARGET = tvm.target.target.micro("host")

#
# 为物理硬件编译
# 在物理硬件上运行时,选择描述硬件的 TARGET 和 BOARD。
# 下面的示例中选择 STM32F746 Nucleo target 和单板。另一种选择是
# STM32F746 Discovery 板。由于该板具有与 Nucleo 相同的 MCU
# 板,但是一些接线和配置不同,选择「stm32f746g_disco」
# 板生成正确的固件镜像。
#

if use_physical_hw:
boards_file = pathlib.Path(tvm.micro.get_microtvm_template_projects("zephyr")) / "boards.json"
with open(boards_file) as f:
boards = json.load(f)

BOARD = os.getenv("TVM_MICRO_BOARD", default="nucleo_f746zg")
TARGET = tvm.target.target.micro(boards[BOARD]["model"])

#
# 对于某些单板,Zephyr 默认使用 QEMU 模拟运行,例如,下面
# TARGET 和 BOARD 用于为 mps2-an521 开发板构建 microTVM 固件。自从那块板
# 在 Zephyr 上默认模拟运行,板名称添加后缀「-qemu」通知
# microTVM 必须使用 QEMU 传输器与开发板通信。如果名称
# 已经有前缀「qemu_」,比如「qemu_x86」,就不需要加上那个后缀了。
#
# TARGET = tvm.target.target.micro("mps2_an521")
# BOARD = "mps2_an521-qemu"

为 target 编译模型:

with tvm.transform.PassContext(
opt_level=3, config={"tir.disable_vectorize": True}, disabled_pass=["AlterOpLayout"]
):
module = relay.build(mod, target=TARGET, runtime=RUNTIME, params=params)

# 检查编译输出
# ---------------------------------
#
# 编译过程产生了一些计算图中实现算子的 C 代码。
# 可以通过打印 CSourceModule 内容来检查它(本教程
# 只打印前 10 行):

c_source_module = module.get_lib().imported_modules[0]
assert c_source_module.type_key == "c", "tutorial is broken"

c_source_code = c_source_module.get_source()
first_few_lines = c_source_code.split("\n")[:10]
assert any(
l.startswith("TVM_DLL int32_t tvmgen_default_") for l in first_few_lines
), f"tutorial is broken: {first_few_lines!r}"
print("\n".join(first_few_lines))

# 编译生成的代码
# ----------------------------
#
# 下面需要将生成的 C 代码合并到一个项目中,以便在
# 设备中运行推理。最简单的方法是自己集成,使用 microTVM 的标准输出格式
# (:doc:`模型库格式` `</dev/model_library_format>`),这是具有标准布局的 tarball:

# 获取可以存储 tarball 的临时路径(作为教程运行)。

fd, model_library_format_tar_path = tempfile.mkstemp()
os.close(fd)
os.unlink(model_library_format_tar_path)
tvm.micro.export_model_library_format(module, model_library_format_tar_path)

with tarfile.open(model_library_format_tar_path, "r:*") as tar_f:
print("\n".join(f" - {m.name}" for m in tar_f.getmembers()))

# 清理:
os.unlink(model_library_format_tar_path)

# TVM 还为嵌入式平台提供了一个标准的方式来自动生成一个独立的
# 项目,编译并烧录到一个 target,使用标准的 TVM RPC 与它通信。
# 模型库格式用作此过程的模型输入。
# 平台为嵌入时提供了集成,可以被 TVM 直接用于主机驱动
# 推理和自动调优。这种集成由
# `microTVM 项目 API` [https://github.com/apache/tvm-rfcs/blob/main/rfcs/0008-microtvm-project-api.md](https://github.com/apache/tvm-rfcs/blob/main/rfcs/0008-microtvm-project-api.md)_提供。
#
# 嵌入式平台需要提供一个包含 microTVM API Server 的模板项目(通常,
# 存在于根目录中的“microtvm_api_server.py”文件中)。本教程使用示例“主机”
# 项目(使用 POSIX 子进程和管道模拟设备):

template_project_path = pathlib.Path(tvm.micro.get_microtvm_template_projects("crt"))
project_options = {} # 可以使用 TVM 提供特定于平台 options。

# 编译物理硬件(或仿真板,如 mps_an521)
# --------------------------------------------------------------------------
# 对于物理硬件,可以通过使用不同的模板项目来试用 Zephyr 平台
# 和选项:
#

if use_physical_hw:
template_project_path = pathlib.Path(tvm.micro.get_microtvm_template_projects("zephyr"))
project_options = {"project_type": "host_driven", "zephyr_board": BOARD}

# 创建临时目录
temp_dir = tvm.contrib.utils.tempdir()
generated_project_dir = temp_dir / "generated-project"
generated_project = tvm.micro.generate_project(
template_project_path, module, generated_project_dir, project_options
)

# 构建并刷新项目
generated_project.build()
generated_project.flash()

输出结果:

// tvm target: c -keys=cpu -link-params=0 -model=host
#define TVM_EXPORTS
#include "tvm/runtime/c_runtime_api.h"
#include "tvm/runtime/c_backend_api.h"
#include <math.h>
#ifdef __cplusplus
extern "C"
#endif
TVM_DLL int32_t tvmgen_default_fused_nn_dense_add(void* args, int32_t* arg_type_ids, int32_t num_args, void* out_ret_value, int32_t* out_ret_tcode, void* resource_handle) {
void* arg_placeholder = (((TVMValue*)args)[0].v_handle);
- .
- ./codegen
- ./codegen/host
- ./codegen/host/src
- ./codegen/host/src/default_lib0.c
- ./codegen/host/src/default_lib1.c
- ./codegen/host/src/default_lib2.c
- ./executor-config
- ./executor-config/graph
- ./executor-config/graph/default.graph
- ./metadata.json
- ./parameters
- ./parameters/default.params
- ./src
- ./src/default.relay

接下来,与模拟设备建立 session,并运行计算。with session 这一行通常会刷新连接的微控制器,但在本教程中,它只启动一个子进程来代替连接的微控制器。

with tvm.micro.Session(transport_context_manager=generated_project.transport()) as session:
graph_mod = tvm.micro.create_local_graph_executor(
module.get_graph_json(), session.get_system_lib(), session.device
)

# 使用「relay.build」产生的降级参数设置模型参数。
graph_mod.set_input(**module.get_params())

# 模型使用单个 float32 值,并返回预测的正弦值。
# 为传递输入值,我们构造一个带有单个构造值的 tvm.nd.array 对象作为输入。
# 这个模型可接收的输入为 0 到 2Pi。
graph_mod.set_input(input_tensor, tvm.nd.array(np.array([0.5], dtype="float32")))
graph_mod.run()

tvm_output = graph_mod.get_output(0).numpy()
print("result is: " + str(tvm_output))

输出结果:

result is: [[0.4443792]]

下载 Python 源代码:micro_tflite.py

下载 Jupyter Notebook:micro_tflite.ipynb