Apache TVM 是一個深度的深度學習編譯框架，適用于 CPU、GPU 和各種機器學習加速芯片。更多 TVM 中文文檔可訪問 →https://tvm.hyper.ai/

作者：Mehrdad Hessar

本教程展示了如何使用 microTVM 構建 MLPerfTiny 提交。該教程演示了從 MLPerfTiny 基準模型中導入一個 TFLite 模型，使用 TVM 進行編譯，并生成一個可以刷寫到支持 Zephyr 的板上的 Zephyr 項目，以使用 EEMBC runner 對模型進行基準測試的步驟。

安裝 microTVM Python 依賴項

TVM 不包含用于 Python 串行通信的軟件包，因此在使用 microTVM 之前，我們必須安裝它。我們還需要 TFLite 來加載模型。

pip install pyserial==3.5 tflite==2.1

import os
import pathlib
import tarfile
import tempfile
import shutil

安裝 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"

注意：僅在您有意使用 CMSIS-NN 代碼生成器生成此提交時安裝 CMSIS-NN。

安裝 Install CMSIS-NN

CMSIS_SHA="51263182d16c92649a48144ba56c0945f9fce60e"
CMSIS_URL="http://github.com/ARM-software/CMSIS_5/archive/${CMSIS_SHA}.tar.gz"
export CMSIS_PATH=/content/cmsis
DOWNLOAD_PATH="/content/${CMSIS_SHA}.tar.gz"
mkdir ${CMSIS_PATH}
wget ${CMSIS_URL} -O "${DOWNLOAD_PATH}"
tar -xf "${DOWNLOAD_PATH}" -C ${CMSIS_PATH} --strip-components=1
rm ${DOWNLOAD_PATH}

CMSIS_NN_TAG="v4.0.0"
CMSIS_NN_URL="https://github.com/ARM-software/CMSIS-NN.git"
git clone ${CMSIS_NN_URL} --branch ${CMSIS_NN_TAG} --single-branch ${CMSIS_PATH}/CMSIS-NN

導入 Python 依賴

import tensorflow as tf
import numpy as np

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
import tvm.micro.testing
from tvm.micro.testing.utils import (
    create_header_file,
    mlf_extract_workspace_size_bytes,
)

導入 Visual Wake Word Model

首先，從 MLPerfTiny 下載并導入 Visual Wake Word (VWW) TFLite 模型。該模型最初來自 MLPerf Tiny 倉庫。我們還捕獲了來自 TFLite 模型的元數據信息，如輸入/輸出名稱、量化參數等，這些信息將在接下來的步驟中使用。

我們使用索引來構建各種模型的提交。索引定義如下：要構建另一個模型，您需要更新模型 URL、簡短名稱和索引號。

關鍵詞識別（KWS）1

視覺喚醒詞（VWW）2

異常檢測（AD）3

圖像分類（IC）4

如果您想要使用 CMSIS-NN 構建提交，請修改 USE_CMSIS 環境變量。

export USE_CMSIS=1

MODEL_URL = "https://github.com/mlcommons/tiny/raw/bceb91c5ad2e2deb295547d81505721d3a87d578/benchmark/training/visual_wake_words/trained_models/vww_96_int8.tflite"
MODEL_PATH = download_testdata(MODEL_URL, "vww_96_int8.tflite", module="model")

MODEL_SHORT_NAME = "VWW"
MODEL_INDEX = 2

USE_CMSIS = os.environ.get("TVM_USE_CMSIS", False)

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)

interpreter = tf.lite.Interpreter(model_path=str(MODEL_PATH))
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

input_name = input_details[0]["name"]
input_shape = tuple(input_details[0]["shape"])
input_dtype = np.dtype(input_details[0]["dtype"]).name
output_name = output_details[0]["name"]
output_shape = tuple(output_details[0]["shape"])
output_dtype = np.dtype(output_details[0]["dtype"]).name

# 從 TFLite 模型中提取量化信息。
# 除了異常檢測模型外，所有其他模型都需要這樣做，
# 因為對于其他模型，我們從主機發送量化數據到解釋器，
# 然而，對于異常檢測模型，我們發送浮點數據，量化信息
# 在微控制器上進行。
if MODEL_SHORT_NAME != "AD":
    quant_output_scale = output_details[0]["quantization_parameters"]["scales"][0]
    quant_output_zero_point = output_details[0]["quantization_parameters"]["zero_points"][0]

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

定義目標、運行時和執行器

現在我們需要定義目標、運行時和執行器來編譯這個模型。在本教程中，我們使用預先編譯（Ahead-of-Time，AoT）進行編譯，并構建一個獨立的項目。這與使用主機驅動模式的 AoT 不同，其中目標會使用主機驅動的 AoT 執行器與主機通信以運行推理。

# 使用 C 運行時 (crt)
RUNTIME = Runtime("crt")

# 使用帶有 `unpacked-api=True` 和 `interface-api=c` 的 AoT 執行器。`interface-api=c` 強制
# 編譯器生成 C 類型的函數 API，而 `unpacked-api=True` 強制編譯器生成最小的未打包格式輸入，
# 這減少了調用模型推理層時的堆棧內存使用。
EXECUTOR = Executor(
    "aot",
    {"unpacked-api": True, "interface-api": "c", "workspace-byte-alignment": 8},
)

# 選擇一個 Zephyr 板
BOARD = os.getenv("TVM_MICRO_BOARD", default="nucleo_l4r5zi")

# 使用 BOARD 獲取完整的目標描述
TARGET = tvm.micro.testing.get_target("zephyr", BOARD)

編譯模型并導出模型庫格式

現在，我們為目標編譯模型。然后，我們為編譯后的模型生成模型庫格式。我們還需要計算編譯后的模型所需的工作空間大小。

config = {"tir.disable_vectorize": True}
if USE_CMSIS:
    from tvm.relay.op.contrib import cmsisnn

    config["relay.ext.cmsisnn.options"] = {"mcpu": TARGET.mcpu}
    relay_mod = cmsisnn.partition_for_cmsisnn(relay_mod, params, mcpu=TARGET.mcpu)

with tvm.transform.PassContext(opt_level=3, config=config):
    module = tvm.relay.build(
        relay_mod, target=TARGET, params=params, runtime=RUNTIME, executor=EXECUTOR
    )

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

生成輸入/輸出頭文件

為了使用 AoT 創建 microTVM 獨立項目，我們需要生成輸入和輸出頭文件。這些頭文件用于將生成的代碼中的輸入和輸出 API 與獨立項目的其余部分連接起來。對于此特定提交，我們只需要生成輸出頭文件，因為輸入 API 調用是以不同的方式處理的。

extra_tar_dir = tvm.contrib.utils.tempdir()
extra_tar_file = extra_tar_dir / "extra.tar"

with tarfile.open(extra_tar_file, "w:gz") as tf:
    create_header_file(
        "output_data",
        np.zeros(
            shape=output_shape,
            dtype=output_dtype,
        ),
        "include/tvm",
        tf,
    )

創建項目、構建并準備項目 tar 文件

現在我們有了編譯后的模型作為模型庫格式，可以使用 Zephyr 模板項目生成完整的項目。首先，我們準備項目選項，然后構建項目。最后，我們清理臨時文件并將提交項目移動到當前工作目錄，可以在開發套件上下載并使用。

input_total_size = 1
for i in range(len(input_shape)):
    input_total_size *= input_shape[i]

template_project_path = pathlib.Path(tvm.micro.get_microtvm_template_projects("zephyr"))
project_options = {
    "extra_files_tar": str(extra_tar_file),
    "project_type": "mlperftiny",
    "board": BOARD,
    "compile_definitions": [
        f"-DWORKSPACE_SIZE={workspace_size + 512}",  # Memory workspace size, 512 is a temporary offset
        # since the memory calculation is not accurate.
        f"-DTARGET_MODEL={MODEL_INDEX}",  # Sets the model index for project compilation.
        f"-DTH_MODEL_VERSION=EE_MODEL_VERSION_{MODEL_SHORT_NAME}01",  # Sets model version. This is required by MLPerfTiny API.
        f"-DMAX_DB_INPUT_SIZE={input_total_size}",  # Max size of the input data array.
    ],
}

if MODEL_SHORT_NAME != "AD":
    project_options["compile_definitions"].append(f"-DOUT_QUANT_SCALE={quant_output_scale}")
    project_options["compile_definitions"].append(f"-DOUT_QUANT_ZERO={quant_output_zero_point}")

if USE_CMSIS:
    project_options["compile_definitions"].append(f"-DCOMPILE_WITH_CMSISNN=1")

# 注意：根據您使用的板子可能需要調整這個值。
project_options["config_main_stack_size"] = 4000

if USE_CMSIS:
    project_options["cmsis_path"] = os.environ.get("CMSIS_PATH", "/content/cmsis")

generated_project_dir = temp_dir / "project"

project = tvm.micro.project.generate_project_from_mlf(
    template_project_path, generated_project_dir, model_tar_path, project_options
)
project.build()

# 清理構建目錄和額外的工件
shutil.rmtree(generated_project_dir / "build")
(generated_project_dir / "model.tar").unlink()

project_tar_path = pathlib.Path(os.getcwd()) / "project.tar"
with tarfile.open(project_tar_path, "w:tar") as tar:
    tar.add(generated_project_dir, arcname=os.path.basename("project"))

print(f"The generated project is located here: {project_tar_path}")

使用此項目與您的板子

既然我們有了生成的項目，您可以在本地使用該項目將板子刷寫并準備好運行 EEMBC runner 軟件。要執行此操作，請按照以下步驟操作：

tar -xf project.tar
cd project
mkdir build
cmake ..
make -j2
west flash

現在，您可以按照這些說明將您的板子連接到 EEMBC runner 并在您的板子上對此模型進行基準測試。

下載 Python 源代碼：micro_mlperftiny.py

下載 Jupyter notebook：micro_mlperftiny.ipynb