😶‍🌫️

M5Stack Module LLM で カスタムモデルを動かす ~ 2.python用runtimeビルド

2024/12/10に公開

はじめに

この記事では@PINTO03091氏によるYOLOv9カスタムモデルのwoholebody17のONNXモデルをM5Stack Module LLMのNPUで実行するためのruntimeをビルドします。

プロジェクトの概要については0.概要をご確認ください。

環境

runtimeのビルドはM5Stack LLM Module上で行います。
LANに接続してSSHログインすると作業がしやすいためDebugボードを使用してEther接続します。
オンボードのeMMC上でも作業可能ですが、ストレージに限りがあるためここではMicroSDカード上にさくせいしたワークスペースで実行します。

M5Stack LLM Moduleの接続

Module LLMのビスを4本取り外と、樹脂フレームが外せます。さらに黒いスピーカーが両面テープで固定されているのでこれを外すとフラットケーブルのコネクタが見えます。デバグボードに付属するフラットケーブルを使って、デバグボードとModule LLMをつなぎます。
つないだ後は部品をなくさない様に、樹脂フレームとスピーカーを取り付けておいてください。

ケーブル類のつなぎこみは、以下の3点です。
①Ethernet接続用のLANケーブルをデバグボードに接続
②Serial通信用にPCとデバグポートをTypeC USBケーブルで接続
③電源供給用にモジュールのTypeC USBポートにPCやモバイルバッテリーから接続

M5Stack LLM Moduleへのログイン

M5Stack LLM Moduleへのログイン方法はいくつかありますが、もっとも簡単なのがデバグボードのUSB Serialを通じたログインかと思います。
ケーブル類の接続ができたら、ターミナルソフトなどを使ってCOMポートを通じて接続します。
デバイスマネージャから確認するとM5Stack LLM ModuleはUSB-SERIAL CH340として認識しているため、COMポート番号が確認できます。

USB Serialからはパスワードなしでログインできます。ログインできたらip aコマンドでipアドレスを表示させてSSHでログインします。USB Serialでも作業は進められますが、VSCodeからSSHでリモートログインする方が作業が簡単です。
私の環境ではIPアドレスは192.168.50.50です。

/# ip a
1: lo: <LOOPBACK,UP,LOWER_UP> mtu 65536 qdisc noqueue state UNKNOWN group default qlen 1000
    link/loopback 00:00:00:00:00:00 brd 00:00:00:00:00:00
    inet 127.0.0.1/8 scope host lo
       valid_lft forever preferred_lft forever
    inet6 ::1/128 scope host 
       valid_lft forever preferred_lft forever
2: eth0: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1500 qdisc mq state UP group default qlen 1000
    link/ether 2e:73:f9:a2:38:e1 brd ff:ff:ff:ff:ff:ff
    inet 192.168.50.50/24 brd 192.168.50.255 scope global eth0
       valid_lft forever preferred_lft forever
    inet6 ::7c9a:d328:2c73:f9ff:fea2:38e1/64 scope global dynamic mngtmpaddr 
       valid_lft 595sec preferred_lft 595sec
    inet6 fe80::2c73:f9ff:fea2:38e1/64 scope link 
       valid_lft forever preferred_lft forever
3: sit0@NONE: <NOARP> mtu 1480 qdisc noop state DOWN group default qlen 1000
    link/sit 0.0.0.0 brd 0.0.0.0

Ethernet接続は接続のたびにIPアドレスが変わってしまうため、固定IPの設定をした方が便利です。
/etc/network/interfacesを次の内容を参考に設定すると固定IP化できます。
address、netmask、gateway、dns-namesrversはそれぞれの環境に合わせてください。

source /etc/network/interfaces.d/*

allow-hotplug eth0
#iface eth0 inet dhcp
iface eth0 inet static
    address 192.168.50.50
    netmask 255.255.255.0
    gateway 192.168.50.1
    dns-nameservers 192.168.50.1

IPアドレスがわかったらVSCodeなどからSSHログインしてください。
デフォルトのアカウントは
user:root
pass:123456
です。

作業環境を整える

クロスコンパイルも可能ですが、環境づくりが面倒なのでModule上で行います。
作業スペース確保のため、ext4フォーマットしたMicroSDを用意してこの中を作業スペースとします。
今回は16GBのMicroSDを使用しましたが、8GB以上あれば作業可能です。
ストレージ容量面ではeMMC上の/optでも作業は可能で処理も早く済むかと思いますが、あまり余裕がないため多くのモデルを保存すると容量不足になる可能性があります。

MicroSDを挿入すると、
/mnt/mmcblk1p1/
にマウントされますのでここに作業スペースを作ります。

CD /mnt/mmcb1k1p1
mkdir /workspace

必要なツールをインストールします。

apt update
apt install unzip
apt install cmake
git clone https://github.com/AXERA-TECH/ax-samples
cd ax-samples
mkdir -p ./3rdparty
wget https://github.com/AXERA-TECH/ax-samples/releases/download/v0.1/opencv-aarch64-linux-gnu-gcc-7.5.0.zip
unzip opencv-aarch64-linux-gnu-gcc-7.5.0.zip -d ./3rdparty
rm opencv-aarch64-linux-gnu-gcc-7.5.0.zip
mkdir build
cd build
git clone https://github.com/AXERA-TECH/ax620q_bsp_sdk.git
export ax_bsp=$PWD/ax620q_bsp_sdk/msp/out/arm64_glibc/

以上でランタイムのビルドに必要なツール類がそろいました。
/mnt/mmcb1k1p1/workspace/ax_samples/examples/ax620e
には、多数のCVモデル用の実行ファイルがあります。すでに準備されている設定を利用すると、これらの実行ファイルをビルド出来ますので試しにビルドする際は以下のコマンドで実行できます。

サンプルコードをビルドする

cd /mnt/mmcblk1p1/workspace/ax-samples/build
cmake -DBSP_MSP_DIR=${ax_bsp}/ -DAXERA_TARGET_CHIP=ax630c ..
make
make install

完成した実行ファイルは、
/mnt/mmcb1k1p1/workspace/ax_samples/build/install/ax630c
に保存されています。

YOLOv9のpython用runtimeをビルドする

続いて、python用のモジュールとしてビルドします。
pythonモジュールへの改編にはpybindを使用するためモジュールをインストールします。

pip install pybind11

/mnt/mmcb1k1p1/workspace/ax_samples/examples/ax620e/ax_yolov9_module_steps.cc
として以下のソースを保存します。

/*
* AXERA is pleased to support the open source community by making ax-samples available.
*
* Copyright (c) 2022, AXERA Semiconductor (Shanghai) Co., Ltd. All rights reserved.
*
* Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
* in compliance with the License. You may obtain a copy of the License at
*
* https://opensource.org/licenses/BSD-3-Clause
*
* Unless required by applicable law or agreed to in writing, software distributed
* under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
* CONDITIONS OF ANY KIND, either express or implied. See the License for the
* specific language governing permissions and limitations under the License.
*/

/*
* Author: ZHEQIUSHUI
*/
/*
* 改編: airpocket
*/


#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/numpy.h>
#include <cstdio>
#include <cstring>
#include <numeric>
#include <opencv2/opencv.hpp>
#include "base/common.hpp"
#include "base/detection.hpp"
#include "middleware/io.hpp"
#include "utilities/args.hpp"
#include "utilities/cmdline.hpp"
#include "utilities/file.hpp"
#include "utilities/timer.hpp"
#include <ax_sys_api.h>
#include <ax_engine_api.h>

const int DEFAULT_IMG_H = 640;
const int DEFAULT_IMG_W = 640;

int NUM_CLASS = 80;
const int DEFAULT_LOOP_COUNT = 1;
const float PROB_THRESHOLD = 0.05f;
const float NMS_THRESHOLD = 0.80f;

namespace py = pybind11;

namespace ax {

    // 共通のオブジェクトを静的変数で管理
    AX_ENGINE_HANDLE handle = nullptr;
    AX_ENGINE_IO_INFO_T* io_info;
    AX_ENGINE_IO_T io_data;


    // 検出結果をPythonに返す形式
    std::vector<std::tuple<int, float, std::array<int, 4>>> post_process(
        AX_ENGINE_IO_INFO_T* io_info, AX_ENGINE_IO_T* io_data, const cv::Mat& mat,
        int input_w, int input_h) {

        std::vector<detection::Object> proposals;
        std::vector<detection::Object> objects;
        //timer timer_postprocess;

        for (int i = 0; i < 3; ++i) {
            auto feat_ptr = (float*)io_data->pOutputs[i].pVirAddr;
            int32_t stride = (1 << i) * 8;
            detection::generate_proposals_yolov9(stride, feat_ptr, PROB_THRESHOLD, proposals, input_w, input_h, NUM_CLASS);
        }

        detection::get_out_bbox(proposals, objects, NMS_THRESHOLD, input_h, input_w, mat.rows, mat.cols);

        // Python に返すリスト
        std::vector<std::tuple<int, float, std::array<int, 4>>> result;

        for (const auto& obj : objects) {
            int x1 = obj.rect.x;
            int y1 = obj.rect.y;
            int x2 = obj.rect.x + obj.rect.width;
            int y2 = obj.rect.y + obj.rect.height;

            //確認のためデバッグ出力
            //fprintf(stdout, " %d:  %.2f%%, [%4d, %4d, %4d, %4d]\n",
            //        obj.label, obj.prob * 100, x1, y1, x2, y2);

            // クラス名は追加しない
            result.emplace_back(
                obj.label, obj.prob,
                std::array<int, 4>{x1, y1, x2, y2});
            
        }
        return result;
        //try {
        //    detection::draw_objects(mat, objects, CLASS_NAMES, "yolov9_out");
        //    fprintf(stdout, "Post-process drawing completed successfully.\n");
        //} catch (const std::exception& e) {
        //    fprintf(stderr, "Error in post_process: %s\n", e.what());
        //    throw;
        //}
    }

    int initialize(
        const std::string& model, 
        int num_class) {
    
        // AXERAエンジンの初期化
        AX_SYS_Init();

        // 1. NPU 初期化
        AX_ENGINE_NPU_ATTR_T npu_attr;
        memset(&npu_attr, 0, sizeof(npu_attr));
        npu_attr.eHardMode = AX_ENGINE_VIRTUAL_NPU_DISABLE;
        auto ret = AX_ENGINE_Init(&npu_attr);
        if (ret != 0) {
            fprintf(stderr, "Error: AX_ENGINE_Init failed with error code %d\n", ret);
            throw std::runtime_error("NPU initialization failed");
        }
        fprintf(stdout, "AX_ENGINE_Init succeeded.\n");

        // 2. モデルロード
        std::vector<char> model_buffer;
        if (!utilities::read_file(model, model_buffer)) {
            fprintf(stderr, "Error: Failed to read model file: %s\n", model.c_str());
            AX_ENGINE_Deinit();
            throw std::runtime_error("Model loading failed");
        }
        fprintf(stdout, "Model file loaded successfully. Size: %zu bytes\n", model_buffer.size());

        // 3. ハンドル作成
        ret = AX_ENGINE_CreateHandle(&handle, model_buffer.data(), model_buffer.size());
        if (ret != 0) {
            fprintf(stderr, "Error: AX_ENGINE_CreateHandle failed with error code %d\n", ret);
            AX_ENGINE_Deinit();
            throw std::runtime_error("Handle creation failed");
        }
        fprintf(stdout, "Engine handle created successfully.\n");

        // 4. コンテキスト作成
        ret = AX_ENGINE_CreateContext(handle);
        if (ret != 0) {
            fprintf(stderr, "Error: AX_ENGINE_CreateContext failed with error code %d\n", ret);
            AX_ENGINE_DestroyHandle(handle);
            AX_ENGINE_Deinit();
            throw std::runtime_error("Context creation failed");
        }
        fprintf(stdout, "Engine context created successfully.\n");

        // 5. IO 情報取得
        ret = AX_ENGINE_GetIOInfo(handle, &io_info);
        if (ret != 0) {
            fprintf(stderr, "Error: AX_ENGINE_GetIOInfo failed with error code %d\n", ret);
            AX_ENGINE_DestroyHandle(handle);
            AX_ENGINE_Deinit();
            throw std::runtime_error("IO information retrieval failed");
        }
        fprintf(stdout, "Engine IO info retrieved successfully.\n");

        // 6. IO メモリ確保
        ret = middleware::prepare_io(io_info, &io_data, std::make_pair(AX_ENGINE_ABST_DEFAULT, AX_ENGINE_ABST_CACHED));
        if (ret != 0) {
            fprintf(stderr, "Error: prepare_io failed with error code %d\n", ret);
            AX_ENGINE_DestroyHandle(handle);
            AX_ENGINE_Deinit();
            throw std::runtime_error("IO memory allocation failed");
        }
        fprintf(stdout, "IO allocated successfully.\n");
        
        return 1;
    }

    // 終了関数
    int finalize() {
        if (!handle) return 0;

        middleware::free_io(&io_data);
        AX_ENGINE_DestroyHandle(handle);
        AX_ENGINE_Deinit();
        handle = nullptr;
        io_info = nullptr;
        return 1;
    }

    std::vector<std::tuple<int, float, std::array<int, 4>>> run_model(
        const std::string& model, 
        const std::vector<uint8_t>& data, 
        cv::Mat& mat, 
        int input_h, 
        int input_w) {

        // 入力データ挿入
        auto ret = middleware::push_input(data, &io_data, io_info);
        if (ret != 0) {
            fprintf(stderr, "Error: push_input failed with error code %d\n", ret);
            middleware::free_io(&io_data);
            AX_ENGINE_DestroyHandle(handle);
            AX_ENGINE_Deinit();
            throw std::runtime_error("Input data insertion failed");
        }

        // モデル実行
        ret = AX_ENGINE_RunSync(handle, &io_data);
        if (ret != 0) {
            fprintf(stderr, "Error: AX_ENGINE_RunSync failed with error code %d\n", ret);
            middleware::free_io(&io_data);
            AX_ENGINE_DestroyHandle(handle);
            AX_ENGINE_Deinit();
            throw std::runtime_error("Model execution failed");
        }
        // 結果処理
        auto results = post_process(io_info, &io_data, mat, input_w, input_h);
        return results;
    }

}

std::vector<std::tuple<int, float, std::array<int, 4>>> inference(
    const std::string& model_file, 
    py::array_t<uint8_t> image_array, 
    std::array<int, 2> input_size, 
    int num_class) {

    // グローバル変数を更新
    NUM_CLASS = num_class;

    // NumPy配列からOpenCVのMatに変換
    py::buffer_info buf = image_array.request();
    if (buf.ndim != 3 || buf.shape[2] != 3) {
        throw std::runtime_error("Input image must be a 3-dimensional array with 3 channels.");
    }
    cv::Mat mat(buf.shape[0], buf.shape[1], CV_8UC3, buf.ptr);

    // cv::Matからstd::vector<uint8_t>への変換
    std::vector<uint8_t> image_data(mat.total() * mat.elemSize());
    std::memcpy(image_data.data(), mat.data, image_data.size());

    // run_modelの呼び出し
    auto results = ax::run_model(model_file, image_data, mat, input_size[0], input_size[1]);
    return results;
}

PYBIND11_MODULE(ax_yolov9_module, m) {
    m.doc() = "AXERA YOLOv9 Python bindings";

    // 初期化関数
    m.def("initialize", &ax::initialize, "Initialize the AXERA engine",
          py::arg("model_file"), 
          py::arg("num_class") = 80);

    m.def("inference", &inference, "Inference",
          py::arg("model_file"),
          py::arg("image_array"),
          py::arg("input_size") = std::array<int, 2>{640, 640},
          py::arg("num_class") = 80);

    // 終了関数
    m.def("finalize", &ax::finalize, "Finalize the AXERA engine");
}

また、/mnt/mmcb1k1p1/workspace/ax_samples/examples/ax620e/CmakeLists.txtを以下の通り修正します。

# Set library directories
link_directories("/mnt/mmcblk1p1/workspace/ax-samples/build/ax620q_bsp_sdk/msp/out/arm64_glibc/lib")

# Add ax_yolov9_py_module as a Python module
add_library(ax_yolov9_module MODULE ax_yolov9_module_steps.cc)

# Locate pybind11
list(APPEND CMAKE_PREFIX_PATH "/usr/local/lib/python3.10/dist-packages/pybind11/share/cmake/pybind11")
find_package(pybind11 REQUIRED)

# Set module properties
set_target_properties(ax_yolov9_module PROPERTIES PREFIX "" SUFFIX ".so")

# Include necessary directories
include_directories(${CMAKE_SOURCE_DIR}/include)
include_directories(/usr/include/python3.10)
include_directories(/usr/local/lib/python3.10/dist-packages/pybind11)
include_directories(/usr/local/lib/python3.10/dist-packages/pybind11/include)
include_directories("/mnt/mmcblk1p1/workspace/ax-samples/build/ax620q_bsp_sdk/msp/out/arm64_glibc/include")


# Link required libraries
target_link_libraries(ax_yolov9_module PRIVATE 
    ${OpenCV_LIBS} 
    /usr/lib/aarch64-linux-gnu/libpython3.10.so 
    ax_engine 
    ax_sys 
    ax_ivps
)

以上の修正ができたら、以下のコマンドでビルドします。

cd /mnt/mmcblk1p1/workspace/ax-samples/build
cmake -DBSP_MSP_DIR=${ax_bsp}/ -DAXERA_TARGET_CHIP=ax630c ..
make

以上でpython用のモジュールが以下のフォルダにax_yolov9_module.soとして保存されています。
/mnt/mmcblk1p1/workspace/ax-samples/build/examples/ax620e

Discussion