(2) 모델 변환하기

pt, onnx 모델을 rknn 형태로 변환하기

예를 들어 YOLO를 오렌지 파이에서 구동하고 싶고, 가지고 있는 것은 ONNX 파일이라고 하면 우리는 다음과 같은 단계로 ONNX 파일을 rknn 파일로 변환할 수 있습니다.

1. rknn 객체 생성
2. configuration 정의
3. ONNX 모델 load
4. rknn 모델 build
5. rknn 모델 export

Code와 함께 따라하기

rknn_model_zoo yolov5 example

rknn_model_zoo/examples/yolov5/python/convert.py at main · airockchip/rknn_model_zoo

import sys

from rknn.api import RKNN

DATASET_PATH = '../../../datasets/COCO/coco_subset_20.txt'
DEFAULT_RKNN_PATH = '../model/yolov5.rknn'
DEFAULT_QUANT = True

def parse_arg():
    if len(sys.argv) < 3:
        print("Usage: python3 {} onnx_model_path [platform] [dtype(optional)] [output_rknn_path(optional)]".format(sys.argv[0]))
        print("       platform choose from [rk3562,rk3566,rk3568,rk3576,rk3588,rv1103,rv1106,rk1808,rv1109,rv1126]")
        print("       dtype choose from [i8, fp] for [rk3562,rk3566,rk3568,rk3576,rk3588,rv1103,rv1106]")
        print("       dtype choose from [u8, fp] for [rk1808,rv1109,rv1126]")
        exit(1)

    model_path = sys.argv[1]
    platform = sys.argv[2]

    do_quant = DEFAULT_QUANT
    if len(sys.argv) > 3:
        model_type = sys.argv[3]
        if model_type not in ['i8', 'u8', 'fp']:
            print("ERROR: Invalid model type: {}".format(model_type))
            exit(1)
        elif model_type in ['i8', 'u8']:
            do_quant = True
        else:
            do_quant = False

    if len(sys.argv) > 4:
        output_path = sys.argv[4]
    else:
        output_path = DEFAULT_RKNN_PATH

    return model_path, platform, do_quant, output_path

if __name__ == '__main__':
    model_path, platform, do_quant, output_path = parse_arg()

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

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

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

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

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

    # Release
    rknn.release()

이 코드를 보면 rknn 모델을 변환할 때는 input의 평균과 표준편차가 필요하다는 것을 알 수 있습니다. 그 부분은 변환하는 모델에 따라서 다르게 설정해줘야 합니다. 예를 들어서 MiDAS 모델을 변환한다고 하면,
https://github.com/isl-org/MiDaS/blob/454597711a62eabcbf7d1e89f3fb9f569051ac9b/midas/model_loader.py#L57

MiDaS/midas/model_loader.py at 454597711a62eabcbf7d1e89f3fb9f569051ac9b · isl-org/MiDaS

이와 같이 코드레벨에서 평균과 표준편차를 알 수 있는 경우가 많습니다. 
 

Configuration

 

• mean_values, std_values
  

• quant_img_RGB2BGR

• target_platform

• quantized_algorithm

• quantized_method

• optimization_level

• quantized_dtype

• custom_string

• dynamic_input