2. microTVM TFLite 指南

备注

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

作者：Tom Gall

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

安装 microTVM Python 依赖项

TVM 不包含用于 Python 串行通信包，因此在使用 microTVM 之前我们必须先安装一个。我们还需要TFLite来加载模型。

pip install pyserial==3.5 tflite==2.1

import os


# 本指南默认运行在使用 TVM 的 C 运行时的 x86 CPU 上，如果你想
# 在 Zephyr 实机硬件上运行，你必须导入 `TVM_MICRO_USE_HW` 环境
# 变量。此外如果你使用 C 运行时，你可以跳过安装 Zephyr。
# 将花费大约20分钟安装 Zephyr。
use_physical_hw = bool(os.getenv("TVM_MICRO_USE_HW"))

安装 Zephyr

# 安装 west 和 ninja
python3 -m pip install west
apt-get install -y ninja-build

# 安装 ZephyrProject
ZEPHYR_PROJECT_PATH="/content/zephyrproject"
export ZEPHYR_BASE=${ZEPHYR_PROJECT_PATH}/zephyr
west init ${ZEPHYR_PROJECT_PATH}
cd ${ZEPHYR_BASE}
git checkout v3.2-branch
cd ..
west update
west zephyr-export
chmod -R o+w ${ZEPHYR_PROJECT_PATH}

# 安装 Zephyr SDK
cd /content
ZEPHYR_SDK_VERSION="0.15.2"
wget "https://github.com/zephyrproject-rtos/sdk-ng/releases/download/v${ZEPHYR_SDK_VERSION}/zephyr-sdk-${ZEPHYR_SDK_VERSION}_linux-x86_64.tar.gz"
tar xvf "zephyr-sdk-${ZEPHYR_SDK_VERSION}_linux-x86_64.tar.gz"
mv "zephyr-sdk-${ZEPHYR_SDK_VERSION}" zephyr-sdk
rm "zephyr-sdk-${ZEPHYR_SDK_VERSION}_linux-x86_64.tar.gz"

# 安装 python 依赖
python3 -m pip install -r "${ZEPHYR_BASE}/scripts/requirements.txt"

导入 Python 依赖项

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

import tvm
import tvm.micro
import tvm.micro.testing
from tvm import relay
import tvm.contrib.utils
from tvm.micro import export_model_library_format
from tvm.contrib.download import download_testdata

model_url = (
    "https://github.com/tlc-pack/web-data/raw/main/testdata/microTVM/model/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 选择下面的「crt」，选择 C Runtime 作为 RUNTIME ，并选择适当的单板/虚拟机来运行它（Zephyr 将创建基于 BOARD 的正确 QEMU 虚拟机）。

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

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

# 运行于物理硬件时，选择描述对应硬件的 TARGET 和 BOARD。
# 下面的示例选择 STM32L4R5ZI Nucleo target 和 board。你可以改变测试板，
# 只需要使用不同的 Zephyr 支持单板导入`TVM_MICRO_BOARD` 变量

if use_physical_hw:
    BOARD = os.getenv("TVM_MICRO_BOARD", default="nucleo_l4r5zi")
    SERIAL = os.getenv("TVM_MICRO_SERIAL", default=None)
    TARGET = tvm.micro.testing.get_target("zephyr", BOARD)

#
# 对于某些单板，Zephyr 默认使用 QEMU 模拟运行，例如，下面
# TARGET 和 BOARD 用于为 mps2-an521 开发板构建 microTVM 固件。
#
# `mps2_an521 = "mps2_an521"`
# `TARGET = tvm.micro.testing.get_target("zephyr", BOARD)`

为 target 编译模型。如果你不需要执行器将默认使用图执行器。

with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
    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 的标准输出格式化模型库格式。这是标准布局的 tarball。

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

temp_dir = tvm.contrib.utils.tempdir()
model_tar_path = temp_dir / "model.tar"
export_model_library_format(module, model_tar_path)

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

# 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。

# 对于物理硬件，可以通过使用不同的模板项目来试用 Zephyr 平台
# 和选项：

if use_physical_hw:
    template_project_path = pathlib.Path(tvm.micro.get_microtvm_template_projects("zephyr"))
        project_options = {
        "project_type": "host_driven",
        "board": BOARD,
        "serial_number": SERIAL,
        "config_main_stack_size": 4096,
        "zephyr_base": os.getenv("ZEPHYR_BASE", default="/content/zephyrproject/zephyr"),
    }

# 创建临时目录

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