RKNNモデル変換チュートリアル例

1. RKNNモデルへの変換

本書では、.rknn 拡張子のモデル評価および実行をサポートしています。TensorFlow、TensorFlow Lite、Caffe、Darknet、ONNX、PyTorchなど一般的な学習済みモデルは、提供されるRKNN-Toolkit2を用いてRKNNモデルへ変換できます。その他のフレームワークで学習したモデルも、一度ONNX形式へ変換してからRKNNへ変換することができます。

変換処理は、RKNNオブジェクトの作成、SDK環境の初期化、各種 load_* インターフェースによるモデル読み込み、build によるRKNNモデル構築、export_rknn によるモデル出力、最後に release によるリソース解放という流れで行います。

2. モデル変換Demoのダウンロード

YOLOv5モデル変換Demoおよび量子化用データセットをダウンロードし、Ubuntu環境へ展開します。Demoには変換スクリプトを含む yolov5_model_convert と、量子化画像を含む quant_dataset が含まれます。

yolov5_model_convert ダウンロードリンク：

https://dl.dragonwake.com/download/rv1126b/AI/demo/model_convert/yolov5_model_convert.tar.bz2

quant_dataset ダウンロードリンク：

https://dl.dragonwake.com/download/rv1126b/AI/demo/model_convert/quant_dataset.zip

3. モデル変換ツールDocker環境へ入る

以下のコマンドを実行し、作業ディレクトリをDockerコンテナ内の /test にマウントします。/dev/bus/usb:/dev/bus/usb はUSBデバイスをコンテナへ渡すための設定です。

docker run -t -i --privileged -v /dev/bus/usb:/dev/bus/usb -v /data/project/sales/csun/rv1126b/jp/AI/demo/model_convert:/test rknn-toolkit2:2.3.2-cp38 /bin/bash

実行は成功しました。下の画像をご覧ください。

注：「/data/project/sales/csun/rv1126b/jp/AI/demo/model_convert」というパスは、実際のUbuntuホスト使用環境に合わせて変更してください。

4. モデル変換操作説明

モデル変換テストDemoは yolov5_model_convert と quant_dataset で構成されています。yolov5_model_convert には変換用スクリプト、quant_dataset には量子化に使用する画像データが格納されています。

圧縮ファイルを解凍します。

unzip quant_dataset.zip
tar -jxvf yolov5_model_convert.tar.bz2

4.1 モデル変換Demoディレクトリ構成

作業ディレクトリ内で、変換スクリプトと量子化データセットが正しく配置されていることを確認します。

yolov5_model_convert フォルダには、主に以下の内容が含まれます。

ファイル	説明
best.onnx	テスト用のモデル
get_list.py	量子化画像リストを生成するスクリプト
rknn_convert.py	モデル変換スクリプト

4.2 量子化画像リストの生成

Docker環境内でモデル変換作業ディレクトリへ移動し、gen_list.py を実行して量子化画像リストを生成します。

cd /test/yolov5_model_convert
python gen_list.py

gen_list.py を実行して、量子化された画像のリストを生成します。

生成された量子化画像のリストは下記のとおりです。

4.3 ONNXモデルをRKNNモデルへ変換

変換スクリプトを実行し、ONNX形式のYOLOv5モデルをRKNN形式へ変換します。変換時には量子化データセットを指定し、ターゲットプラットフォームや量子化設定に応じてRKNNモデルを生成します。

rknn_convert.py スクリプトは、デフォルトでint8量子化を実行します。スクリプトのコードは以下のとおりです。

import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN

ONNX_MODEL = 'best.onnx'
RKNN_MODEL = './yolov5_mask_rk3576.rknn'
DATASET = './pic_path.txt'

QUANTIZE_ON = True


if __name__ == '__main__':

    # Create RKNN object
    rknn = RKNN(verbose=True)

    if not os.path.exists(ONNX_MODEL):
        print('model not exist')
        exit(-1)

    # pre-process config
    print('--> Config model')
    rknn.config(mean_values=[[0, 0, 0]],
                std_values=[[255, 255, 255]],
                target_platform='rk3576')
    print('done')

    # Load ONNX model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL)
    if ret != 0:
        print('Load yolov5 failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)
    if ret != 0:
        print('Build yolov5 failed!')
        exit(ret)
    print('done')

    # Export RKNN model
    print('--> Export RKNN model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export yolov5rknn failed!')
        exit(ret)
    print('done')

ONNXモデル best.onnx を yolov5_model_convert ディレクトリに配置し、rknn_convert.py スクリプトを実行してモデル変換を行います。

python rknn_convert.py

生成されたモデルは下図に示すとおりです。このモデルは、RKNN環境と基板環境の両方で実行可能です。