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

作者：Mehrdad Hessar

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

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

導(dǎo)入 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,
)

導(dǎo)入 Visual Wake Word Model

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

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

關(guān)鍵詞識別（KWS）1

視覺喚醒詞（VWW）2

異常檢測（AD）3

圖像分類（IC）4

如果您想要使用 CMSIS-NN 構(gòu)建提交，請修改 USE_CMSIS 環(huán)境變量。

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 模型中提取量化信息。
# 除了異常檢測模型外，所有其他模型都需要這樣做，
# 因為對于其他模型，我們從主機發(fā)送量化數(shù)據(jù)到解釋器，
# 然而，對于異常檢測模型，我們發(fā)送浮點數(shù)據(jù)，量化信息
# 在微控制器上進行。
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}
)

定義目標(biāo)、運行時和執(zhí)行器

現(xiàn)在我們需要定義目標(biāo)、運行時和執(zhí)行器來編譯這個模型。在本教程中，我們使用預(yù)先編譯（Ahead-of-Time，AoT）進行編譯，并構(gòu)建一個獨立的項目。這與使用主機驅(qū)動模式的 AoT 不同，其中目標(biāo)會使用主機驅(qū)動的 AoT 執(zhí)行器與主機通信以運行推理。

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

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

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

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

編譯模型并導(dǎo)出模型庫格式

現(xiàn)在，我們?yōu)槟繕?biāo)編譯模型。然后，我們?yōu)榫幾g后的模型生成模型庫格式。我們還需要計算編譯后的模型所需的工作空間大小。

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 創(chuàng)建 microTVM 獨立項目，我們需要生成輸入和輸出頭文件。這些頭文件用于將生成的代碼中的輸入和輸出 API 與獨立項目的其余部分連接起來。對于此特定提交，我們只需要生成輸出頭文件，因為輸入 API 調(diào)用是以不同的方式處理的。

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,
    )

創(chuàng)建項目、構(gòu)建并準(zhǔn)備項目 tar 文件

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

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")

# 注意：根據(jù)您使用的板子可能需要調(diào)整這個值。
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()

# 清理構(gòu)建目錄和額外的工件
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}")

使用此項目與您的板子

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

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

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

下載 Python 源代碼：micro_mlperftiny.py

下載 Jupyter notebook：micro_mlperftiny.ipynb