单击 此处 下载完整的示例代码
9. 开发环境中加入 microTVM
本教程描述了将使用 microTVM 编译的模型集成到自定义开发环境所需的步骤。在本教程中,我们使用 STM32CubeIDE 作为目标集成开发环境(IDE),但我们不依赖于此 IDE 的任何特定功能,将 microTVM 集成到其他 IDE 中的步骤类似。在这里,我们还使用了 MLPerf Tiny 的 Visual Wake Word(VWW)模型和 nucleo_l4r5zi 开发板,但相同的步骤也适用于任何其他模型或目标微控制器单元(MCU)。如果您希望在 vww 模型上使用另一个目标 MCU,我们建议选择具有约 512 KB 和约 256 KB 闪存和 RAM 的 Cortex-M4 或 Cortex-M7 设备。
以下是本教程中要执行的步骤的简要概述。
-
首先,我们导入模型,使用 TVM 进行编译,并生成包含模型生成代码以及所有所需 TVM 依赖项的 Model Library Format(MLF)tar 文件。
-
我们还将两个二进制格式的样本图像(一个人和一个非人样本)添加到 .tar 文件中,以用于评估模型。
-
接下来,我们使用 stmCubeMX 生成在 stmCube IDE 中项目的初始化代码。
-
然后 ,我们将我们的 MLF 文件和所需的 CMSIS 库包含到项目中并进行构建。
-
最后,我们烧写设备并在我们的样本图像上评估模型性能。
让我们开始吧。
安装 microTVM Python 依赖项
TVM 不包含用于 Python 串行通信的包,因此在使用 microTVM 之前,我们必须安装一个。我们还需要 TFLite 以加载模型,以及 Pillow 以准备样本图像。
pip install pyserial==3.5 tflite==2.1 Pillow==9.0 typing_extensions
导入 Python 依赖项
如果要在本地运行此脚本,请查看 TVM 在线文档,了解安装 TVM 的说明。
import os
import numpy as np
import pathlib
import json
from PIL import Image
import tarfile
import tvm
from tvm import relay
from tvm.relay.backend import Executor, Runtime
from tvm.contrib.download import download_testdata
from tvm.micro import export_model_library_format
from tvm.relay.op.contrib import cmsisnn
from tvm.micro.testing.utils import create_header_file
导入 TFLite 模型
首先,下载并导 入 Visual Wake Word TFLite 模型。该模型接受一个 96x96x3 的 RGB 图像,并确定图像中是否存在人物。此模型最初来自 MLPerf Tiny 仓库。为了测试该模型,我们使用 COCO 2014 Train images 中的两个样本。
MODEL_URL = "https://github.com/mlcommons/tiny/raw/bceb91c5ad2e2deb295547d81505721d3a87d578/benchmark/training/visual_wake_words/trained_models/vww_96_int8.tflite"
MODEL_NAME = "vww_96_int8.tflite"
MODEL_PATH = download_testdata(MODEL_URL, MODEL_NAME, module="model")
tflite_model_buf = open(MODEL_PATH, "rb").read()
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)
input_shape = (1, 96, 96, 3)
INPUT_NAME = "input_1_int8"
relay_mod, params = relay.frontend.from_tflite(
tflite_model, shape_dict={INPUT_NAME: input_shape}, dtype_dict={INPUT_NAME: "int8"}
)
生成模型库格式文件
首先,我们定义目标、运行时和执行器。然后,我们为目标设备编译模型,最后导出生成的代码和所有必需的依赖项到单个文件中。
# 我们可以使用 TVM 的本地调度或依赖于 CMSIS-NN 内核,使用 TVM 的 Bring-Your-Own-Code (BYOC) 能力。
USE_CMSIS_NN = True
# USMP (Unified Static Memory Planning) 对所有张量进行综合内存规划,以实现最佳内存利用。
DISABLE_USMP = False
# 使用 C 运行时(crt)
RUNTIME = Runtime("crt")
# 我们通过将板名称传递给 `tvm.target.target.micro` 来定义目标。
# 如果您的板型未包含在支持的模型中,您可以定义目标,如下所示:
# TARGET = tvm.target.Target("c -keys=arm_cpu,cpu -mcpu=cortex-m4")
TARGET = tvm.target.target.micro("stm32l4r5zi")
# 使用 AOT 执行器而不是图形或虚拟机执行器。使用未打包的 API 和 C 调用风格。
EXECUTOR = tvm.relay.backend.Executor(
"aot", {"unpacked-api": True, "interface-api": "c", "workspace-byte-alignment": 8}
)
# 现在,我们设置编译配置并为目标编译模型:
config = {"tir.disable_vectorize": True}
if USE_CMSIS_NN:
config["relay.ext.cmsisnn.options"] = {"mcpu": TARGET.mcpu}
if DISABLE_USMP:
config["tir.usmp.enable"] = False
with tvm.transform.PassContext(opt_level=3, config=config):
if USE_CMSIS_NN:
# 当我们使用 CMSIS-NN 时,TVM 在 relay 图中搜索可以转移到 CMSIS-NN 内核的模式。
relay_mod = cmsisnn.partition_for_cmsisnn(relay_mod, params, mcpu=TARGET.mcpu)
lowered = tvm.relay.build(
relay_mod, target=TARGET, params=params, runtime=RUNTIME, executor=EXECUTOR
)
parameter_size = len(tvm.runtime.save_param_dict(lowered.get_params()))
print(f"Model parameter size: {parameter_size}")
# 我们需要选择一个目录来保存我们的文件。
# 如果在 Google Colab 上运行,我们将保存所有内容在 ``/root/tutorial`` 中(也就是 ``~/tutorial``),
# 但是如果在本地运行,您可能希望将其存储在其他位置。
BUILD_DIR = pathlib.Path("/root/tutorial")
BUILD_DIR.mkdir(exist_ok=True)
# 现在,我们将模型导出为一个 tar 文件:
TAR_PATH = pathlib.Path(BUILD_DIR) / "model.tar"
export_model_library_format(lowered, TAR_PATH)
输出:
Model parameter size: 32
PosixPath('/workspace/gallery/how_to/work_with_microtvm/tutorial/model.tar')
将样本图像添加到 MLF 文件中
最后,我们下载两个样本图像(一个人图像和一个非人图像),将它们转换为二进制格式,并存储在两个头文件中。
with tarfile.open(TAR_PATH, mode="a") as tar_file:
SAMPLES_DIR = "samples"
SAMPLE_PERSON_URL = (
"https://github.com/tlc-pack/web-data/raw/main/testdata/microTVM/data/vww_sample_person.jpg"
)
SAMPLE_NOT_PERSON_URL = "https://github.com/tlc-pack/web-data/raw/main/testdata/microTVM/data/vww_sample_not_person.jpg"
SAMPLE_PERSON_PATH = download_testdata(SAMPLE_PERSON_URL, "person.jpg", module=SAMPLES_DIR)
img = Image.open(SAMPLE_PERSON_PATH)
create_header_file("sample_person", np.asarray(img), SAMPLES_DIR, tar_file)
SAMPLE_NOT_PERSON_PATH = download_testdata(
SAMPLE_NOT_PERSON_URL, "not_person.jpg", module=SAMPLES_DIR
)
img = Image.open(SAMPLE_NOT_PERSON_PATH)
create_header_file("sample_not_person", np.asarray(img), SAMPLES_DIR, tar_file)
在这一点上,您已经具备将编译后的模型导入到您的 IDE 并进行评估所需的一切。在 MLF 文件(model.tar)中,您应该找到以下文件层次结构:
/root
├── codegen
├── parameters
├── runtime
├── samples
├── src
├── templates
├── metadata.json
- codegen 文件夹:包含了由 TVM 为您的模型生成的 C 代码。
- runtime 文件夹:包含了目标需要编译生成的 C 代码所需的所有 TVM 依赖项。
- samples 文件夹:包含了用于评估模型的两个生成的样本文件。
- src 文件夹:包含了描述模型的 relay 模块。
- templates 文件夹:包含了两个模板文件,根据您的平台可能需要进行编辑。
- metadata.json 文件:包含有关模型、其层次和内存需求的信息。