[RKNN] rknn-toolkit2

Intro

이번 주제는 지난번에 예고한 rknn-toolkit2 이다. rknn-toolkit2는 AI 모델을 RKNPU가 연산 가능하도록 RKNN 형식의 모델로 변환해주는 프레임워크이다.

이번 글에서 rknn-toolkit2를 이용해 모델을 어떻게 rknn 형식으로 변환하는지 알아보자.

Envrionment

rknn-toolkit2를 사용하기 위해서는 몇 가지 제약사항이 존재한다.

OS는 리눅스만 지원된다.
rknn-toolkit2가 2.3.0 이상일 경우, x86 및 ARM 환경 모두 사용 가능하다.
python 3.6 ~ python 3.12까지 지원한다.
다음과 같은 프레임워크에서 변환된 AI 모델만 지원한다.
- Pytorch, Tensorflow, Tensorflow Lite, ONNX, Caffe, Darknet

Installation

rknn-toolkit2 설치하는 방법은 두 가지이다.

pypi를 통한 설치
github에서 whl 파일을 다운로드하여 수동 설치

각각의 설치방법에 대해 알아보자.

pypi를 이용한 설치

pypi를 이용한 설치는 간단하다. 아래 명령어만 실행하면 설치가 완료된다.

pip install rknn-toolkit2 -i https://pypi.org/simple

whl 파일을 이용한 설치

whl 파일은 rknn-toolkit2의 packages 디렉토리에 존재하며,
사용 환경(x86 or ARM)과 Python 버전에 따라 설치할 파일이 달라진다.

설치환경이 x86(일반 컴퓨터)이라면, x86_64 디렉토리에 들어가 python 버전에 맞는 txt, whl 파일을 다운받으면 된다. ARM 환경이라면 arm64 디렉토리에 들어가면 된다.

txt 파일은 아래와 같이 구분하여 알맞은 환경을 다운로드 한다.

# basic
requirements_cp{python version}-{rknn version}
# example
requirement_cp310-2.3.2.txt # x86, python 3.10
arm64_requirements_cp310_2.3.2.txt # arm, python 3.10

whl 파일은 아래와 같이 구분하여 알맞은 환경을 다운로드 한다.

# basic
rknn_toolkit2-{rknn version}-cp{python version}-cp{python version}-manylinux_2_17_{PC env}.manylinux2014_{PC env}.whl

# example
rknn_toolkit2-2.3.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl # x86, python 3.10
rknn_toolkit2-2.3.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl # arm, python 3.10

txt, whl 파일을 모두 받았다면, 아래의 명령어를 통해 라이브러리 설치를 시작한다.

# example : x86, python 3.10
# 기본 환경 설치
pip install -r requirement_cp310-2.3.2.txt
# rknn-toolkit2 설치
pip install rknn_toolkit2-2.3.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

⚠ Cmake 설치 문제

rknn-toolkit2를 설치하던 도중 아래와 같은 문제가 발생할 수 있다.

Preparing metadata (setup.py) ... error
  error: subprocess-exited-with-error
...
          assert CMAKE, 'Could not find "cmake" executable!'
      AssertionError: Could not find "cmake" executable!
      [end of output]

  note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

이는 cmake 프로그램이 설치되지 않아 생긴 문제이다. cmake 설치 후 rknn-toolkit2를 재설치하면 정상적으로 설치가 완료된다.

sudo apt install cmake
pip install rknn-toolkit2 -i https://pypi.org/simple

RKNN 모델 변환 절차

위 이미지는 RKNPU User Guide V2.3.0 중 3.1절 Model Conversion에 존재하는 이미지 이다. rknn 형식의 모델 변환을 위해서는 다음과 같은 절차가 존재한다.

rknn 객체 초기화 (RKNN init)
변환 구성 설정 (config)
모델 불러오기 (load_model)
모델 구축 (build)
rknn 배포 (export_rknn)
rknn release

위 절차를 토대로 코드를 작성하면 어떻게 될까?

Python Code

위의 절차를 토대로 만든 파이썬 예제 코드이다.

from rknn.api import RKNN

ONNX_MODEL = 'model.onnx'
RKNN_MODEL = 'model.rknn'
QUANTIZE_ON = False

if __name__ == '__main__':

    # RKNN init
    rknn = RKNN(verbose=False, verbose_file='./model_build.log')

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

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

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

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

    # RKNN release
    rknn.release()

이제 각 함수에 파라미터가 무엇이 있는지 알아보자.

RKNN init

rknn = RKNN(verbose=False, verbose_file='')

Params	Type	설명
verbose	bool	로그 정보 표시 여부
verbose_file	str	로그 정보 저장 파일명

Config

rknn.config(
    mean_values=[[0, 0, 0]], 
    std_values=[[255, 255, 255]], 
    target_platform='RK3588'
)

Params	Type	설명
mean_values	list	입력의 평균값
str_values	list	입력의 정규화된 값
target_platform	str	RKNN 플랫폼 (RK3588, RK2118 등)

⚠ 이 외에도 quant_img_RGB2BGR, quantized_algorithm, quantized_dtype 등의 추가 파라미터가 있으며, 자세한 내용은 User Guide의 3.1.2 항목을 참고하세요.

Load Model (onnx 기준)

rknn.load_onnx(model='model.onnx')

Params	Type	설명
model	str	onnx 모델 위치

Build

rknn.build(
    do_quantization=False,
    dataset=''
)

Params	Type	설명
do_quantization	bool	모델 양자화 여부 (양자화 : INT8)
dataset	str	양자화 보정에 사용되는 데이터 파일 (형식 : txt 파일)

⚠ 양자화는 현재 INT8만 지원됩니다. 자세한 내용은 User Guide 6.2.1 항목을 참고하세요.

Export RKNN

rknn.export_rknn(export_path='model.rknn')

Params	Type	설명
export_path	str	rknn 형식 모델 저장명

회고

이번 글에서는 rknn-toolkit2의 기본 사용법과 설치 방법에 대해 정리해 보았다. User Guide의 내용을 많이 작성하고 싶었지만, 아직 직접 실습하지 않은 기능도 많기에 다루지 못한 항목이 많다. 다음에 기회가 된다면 User Guide의 내용을 하나씩 실행해볼까 싶다.

다음에는 변환된 rknn 모델을 실제로 RKNPU에서 실행하는 rknn-toolkit-lite2에 대한 글을 쓰고자 한다.

참고사이트

aircokchip/rknn-toolkit2
GitHub - airockchip/rknn-toolkit2
Contribute to airockchip/rknn-toolkit2 development by creating an account on GitHub.
https://github.com/airockchip/rknn-toolkit2

rknn-toolkit2 User Guide EN v2.3.2
rknn-toolkit2/doc/02_Rockchip_RKNPU_User_Guide_RKNN_SDK_V2.3.2_EN.pdf at master · airockchip/rknn-toolkit2
Contribute to airockchip/rknn-toolkit2 development by creating an account on GitHub.
https://github.com/airockchip/rknn-toolkit2/blob/master/doc/02_Rockchip_RKNPU_User_Guide_RKNN_SDK_V2.3.2_EN.pdf

rknn model zoo (rknn example 모음집)
GitHub - airockchip/rknn_model_zoo
Contribute to airockchip/rknn_model_zoo development by creating an account on GitHub.
https://github.com/airockchip/rknn_model_zoo