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

本教程描述了將使用 microTVM 編譯的模型集成到自定義開發(fā)環(huán)境所需的步驟。在本教程中，我們使用?STM32CubeIDE?作為目標(biāo)集成開發(fā)環(huán)境（IDE），但我們不依賴于此 IDE 的任何特定功能，將 microTVM 集成到其他 IDE 中的步驟類似。在這里，我們還使用了 MLPerf Tiny 的 Visual Wake Word（VWW）模型和 nucleo_l4r5zi 開發(fā)板，但相同的步驟也適用于任何其他模型或目標(biāo)微控制器單元（MCU）。如果您希望在 vww 模型上使用另一個目標(biāo) MCU，我們建議選擇具有約 512 KB 和約 256 KB 閃存和 RAM 的 Cortex-M4 或 Cortex-M7 設(shè)備。

以下是本教程中要執(zhí)行的步驟的簡要概述。

1. 首先，我們導(dǎo)入模型，使用 TVM 進(jìn)行編譯，并生成包含模型生成代碼以及所有所需 TVM 依賴項的?Model Library Format（MLF）tar 文件。
2. 我們還將兩個二進(jìn)制格式的樣本圖像（一個人和一個非人樣本）添加到 .tar 文件中，以用于評估模型。
3. 接下來，我們使用 stmCubeMX 生成在 stmCube IDE 中項目的初始化代碼。
4. 然后，我們將我們的 MLF 文件和所需的 CMSIS 庫包含到項目中并進(jìn)行構(gòu)建。
5. 最后，我們燒寫設(shè)備并在我們的樣本圖像上評估模型性能。

讓我們開始吧。

安裝 microTVM Python 依賴項

TVM 不包含用于 Python 串行通信的包，因此在使用 microTVM 之前，我們必須安裝一個。我們還需要 TFLite 以加載模型，以及 Pillow 以準(zhǔn)備樣本圖像。

pip install pyserial==3.5 tflite==2.1 Pillow==9.0 typing_extensions

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

導(dǎo)入 TFLite 模型

首先，下載并導(dǎo)入 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"}
)

生成模型庫格式文件

首先，我們定義目標(biāo)、運行時和執(zhí)行器。然后，我們?yōu)槟繕?biāo)設(shè)備編譯模型，最后導(dǎo)出生成的代碼和所有必需的依賴項到單個文件中。

# 我們可以使用 TVM 的本地調(diào)度或依賴于 CMSIS-NN 內(nèi)核，使用 TVM 的 Bring-Your-Own-Code (BYOC) 能力。
USE_CMSIS_NN = True

# USMP (Unified Static Memory Planning) 對所有張量進(jìn)行綜合內(nèi)存規(guī)劃，以實現(xiàn)最佳內(nèi)存利用。
DISABLE_USMP = False

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

# 我們通過將板名稱傳遞給 `tvm.target.target.micro` 來定義目標(biāo)。
# 如果您的板型未包含在支持的模型中，您可以定義目標(biāo)，如下所示：
# TARGET = tvm.target.Target("c -keys=arm_cpu,cpu -mcpu=cortex-m4")
TARGET = tvm.target.target.micro("stm32l4r5zi")

# 使用 AOT 執(zhí)行器而不是圖形或虛擬機執(zhí)行器。使用未打包的 API 和 C 調(diào)用風(fēng)格。
EXECUTOR = tvm.relay.backend.Executor(
    "aot", {"unpacked-api": True, "interface-api": "c", "workspace-byte-alignment": 8}
)

# 現(xiàn)在，我們設(shè)置編譯配置并為目標(biāo)編譯模型：
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:
        # 當(dāng)我們使用 CMSIS-NN 時，TVM 在 relay 圖中搜索可以轉(zhuǎn)移到 CMSIS-NN 內(nèi)核的模式。
        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 上運行，我們將保存所有內(nèi)容在 ``/root/tutorial`` 中（也就是 ``~/tutorial``），
# 但是如果在本地運行，您可能希望將其存儲在其他位置。

BUILD_DIR = pathlib.Path("/root/tutorial")

BUILD_DIR.mkdir(exist_ok=True)

# 現(xiàn)在，我們將模型導(dǎo)出為一個 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 文件中?

最后，我們下載兩個樣本圖像（一個人圖像和一個非人圖像），將它們轉(zhuǎn)換為二進(jìn)制格式，并存儲在兩個頭文件中。

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)

在這一點上，您已經(jīng)具備將編譯后的模型導(dǎo)入到您的 IDE 并進(jìn)行評估所需的一切。在 MLF 文件（model.tar）中，您應(yīng)該找到以下文件層次結(jié)構(gòu)：

/root
├── codegen
├── parameters
├── runtime
├── samples
├── src
├── templates
├── metadata.json

• codegen 文件夾：包含了由 TVM 為您的模型生成的 C 代碼。
• runtime 文件夾：包含了目標(biāo)需要編譯生成的 C 代碼所需的所有 TVM 依賴項。
• samples 文件夾：包含了用于評估模型的兩個生成的樣本文件。
• src 文件夾：包含了描述模型的 relay 模塊。
• templates 文件夾：包含了兩個模板文件，根據(jù)您的平臺可能需要進(jìn)行編輯。
• metadata.json 文件：包含有關(guān)模型、其層次和內(nèi)存需求的信息。

生成在您的 IDE 中的項目?

下一步是為目標(biāo)設(shè)備創(chuàng)建一個項目。我們使用 STM32CubeIDE，您可以在此處下載。在本教程中，我們使用的是版本 1.11.0。安裝 STM32CubeIDE 后，請按照以下步驟創(chuàng)建項目：

1. 選擇 File -> New -> STM32Project。目標(biāo)選擇窗口將出現(xiàn)。

2. 轉(zhuǎn)到 “Board Selector” 選項卡，在 “Commercial Part Number” 文本框中鍵入板名稱 “nucleo-l4r5zi”。從右側(cè)顯示的板列表中選擇板，并單擊 “Next”。

3. 輸入項目名稱（例如 microtvm_vww_demo）。我們使用默認(rèn)選項（目標(biāo)語言：C，二進(jìn)制類型：可執(zhí)行文件，項目類型：STM32Cube）。單擊 “Finish”。

4. 一個文本框?qū)⒊霈F(xiàn)，詢問是否要 “以默認(rèn)模式初始化所有外設(shè)？”。點擊 “Yes”。這將生成項目并打開設(shè)備配置工具，您可以使用 GUI 設(shè)置外設(shè)。默認(rèn)情況下啟用了 USB、USART3 和 LPUART1，以及一些 GPIO。

5. 我們將使用 LPUART1 將數(shù)據(jù)發(fā)送到主機 PC。從連接部分中選擇 LPUART1，并將 “Baud Rate” 設(shè)置為 115200，將 “Word Length” 設(shè)置為 8。保存更改并點擊 “Yes” 以重新生成初始化代碼。這應(yīng)該會重新生成代碼并打開您的 main.c 文件。您還可以從左側(cè)的 Project Explorer 面板中找到 main.c，在 microtvm_vww_demo -> Core -> Src 下。

6. 為了進(jìn)行健全性檢查，請復(fù)制下面的代碼并將其粘貼到主函數(shù)的無線循環(huán)（即 While(1) ）部分。

   • 注意：確保您的代碼寫在由 USER CODE BEGIN<…> 和 USER CODE END<…> 包圍的部分內(nèi)。如果重新生成初始化代碼，被包圍之外的代碼將被擦除。

HAL_GPIO_TogglePin(LD2_GPIO_Port, LD2_Pin);
HAL_UART_Transmit(&hlpuart1, "Hello World.\r\n", 14, 100);
HAL_Delay(1000);

7. 從菜單欄中選擇 Project -> Build（或右鍵單擊項目名稱并選擇 Build）。這將構(gòu)建項目并生成 .elf 文件。選擇 Run -> Run 以將二進(jìn)制文件下載到您的 MCU。如果打開了“Edit Configuration”窗口，請直接點擊 “OK”。
8. 在主機機器上打開終端控制臺。在 Mac 上，您可以簡單地使用 “screen <usb_device> 115200” 命令，例如 “screen tty.usbmodemXXXX 115200” 。板上的 LED 應(yīng)該會閃爍，終端控制臺上每秒應(yīng)該會打印出字符串 “Hello World.”。按 “Control-a k” 退出 screen。

將模型導(dǎo)入生成的項目?

要將編譯后的模型集成到生成的項目中，請按照以下步驟操作：

1. 解壓 tar 文件并將其包含在項目中

   • 打開項目屬性（右鍵單擊項目名稱并選擇 “Properties” 或從菜單欄選擇 Project -> Properties）。
   • 選擇 C/C++ General -> Paths and Symbols。選擇 Source Location 選項卡。
   • 如果您將模型解壓縮在項目文件夾內(nèi)，請點擊 “Add Folder” 并選擇 “model” 文件夾（在它出現(xiàn)之前，您可能需要右鍵單擊項目名稱并選擇 “Refresh”）。
   • 如果您在其他地方解壓縮了模型文件，請點擊 “Link Folder” 按鈕，在出現(xiàn)的窗口中選中 “Link to folder in the file system” 復(fù)選框，點擊 “Browse” 并選擇模型文件夾。

2. 如果在編譯模型時使用了 CMSIS-NN，您還需要在項目中包含 CMSIS-NN 源文件。

   • 從?CMSIS-NN 存儲庫下載或克隆文件，并按照上述步驟將 CMSIS-NN 文件夾包含在項目中。

3. 打開項目屬性。在 C/C++ Build -> Settings 中：通過點擊 “+” 按鈕，選擇 “Workspace” ，并導(dǎo)航到以下各個文件夾。將以下文件夾添加到 MCU GCC Compiler 的 Include Paths 列表中（如果是 C++ 項目還需添加到 MCU G++ Compiler 中）：

   • model/runtime/include
   • model/codegen/host/include
   • model/samples
   • CMSIS-NN/Include

4. 從 model/templates 復(fù)制 crt_config.h.template 到 Core/Inc 文件夾，并將其重命名為 crt_config.h。

5. 從 model/templates 復(fù)制 platform.c.template 到 Core/Src 文件夾，并將其重命名為 platform.c。

   • 此文件包含您可能需要根據(jù)平臺編輯的內(nèi)存管理函數(shù)。
   • 在 platform.c 中定義 “TVM_WORKSPACE_SIZE_BYTES” 的值。如果使用 USMP，則只需要比較小的值（例如 1024 字節(jié)）即可。
   • 如果不使用 USMP，請查看 metadata.json 中的 “workspace_size_bytes” 字段以估算所需內(nèi)存。

6. 從構(gòu)建中排除以下文件夾（右鍵單擊文件夾名稱，選擇 Resource Configuration → Exclude from build）。檢查 Debug 和 Release 配置。

   • CMSIS_NN/Tests

7. 從?CMSIS Version 5 存儲庫下載 CMSIS 驅(qū)動程序。

   • 在項目目錄中，刪除 Drivers/CMSIS/Include 文件夾（這是 CMSIS 驅(qū)動程序的舊版本），并將您從下載的版本中復(fù)制的 CMSIS/Core/Include 粘貼到相同位置。

8. 編輯 main.c 文件：

• 包含下列頭文件

#include <stdio.h>
#include <string.h>
#include <stdarg.h>
#include "tvmgen_default.h"
#include "sample_person.h"
#include "sample_not_person.h"

• 在 main 函數(shù)的無限循環(huán)前復(fù)制下面這段代碼。該代碼設(shè)置模型的輸入和輸出

TVMPlatformInitialize();
signed char output[2];
struct tvmgen_default_inputs inputs = {
.input_1_int8 = (void*)&sample_person,
};
struct tvmgen_default_outputs outputs = {
.Identity_int8 = (void*)&output,
};
char msg[] = "Evaluating VWW model using microTVM:\r\n";
HAL_UART_Transmit(&hlpuart1, msg, strlen(msg), 100);
uint8_t sample = 0;
uint32_t timer_val;
char buf[50];
uint16_t buf_len;

• 將以下代碼復(fù)制到無限循環(huán)中。該代碼將在圖片上運行推斷并在控制臺打印結(jié)果。

if (sample == 0)
    inputs.input_1_int8 = (void*)&sample_person;
else
    inputs.input_1_int8 = (void*)&sample_not_person;

timer_val = HAL_GetTick();
tvmgen_default_run(&inputs, &outputs);
timer_val = HAL_GetTick() - timer_val;
if (output[0] > output[1])
    buf_len = sprintf(buf, "Person not detected, inference time = %lu ms\r\n", timer_val);
else
    buf_len = sprintf(buf, "Person detected, inference time = %lu ms\r\n", timer_val);
HAL_UART_Transmit(&hlpuart1, buf, buf_len, 100);

sample++;
if (sample == 2)
    sample = 0;

• 在 main 中定義 TVMLogf 函數(shù)，接受 TVM 運行時在控制臺的報錯

void TVMLogf(const char* msg, ...) {
  char buffer[128];
  int size;
  va_list args;
  va_start(args, msg);
  size = TVMPlatformFormatMessage(buffer, 128, msg, args);
  va_end(args);
  HAL_UART_Transmit(&hlpuart1, buffer, size, 100);
}

9. 在項目屬性中，找到 C/C++ Build -> Settings, MCU GCC Compiler -> Optimization，設(shè)置 Optimization 為 Optimize more (-O2)。

評估模型?

現(xiàn)在，選擇菜單欄中的 Run -> Run 來刷寫 MCU 并運行項目。您應(yīng)該看到 LED 在閃爍，并且控制臺上在打印推理結(jié)果。

下載 Python 源代碼：micro_custom_ide.py

下載 Jupyter notebook：micro_custom_ide.ipynb