[Torch2TFLite] Torch 모델 TFLite 변환 (feat. yolov5)

AI 개발자라면 가장 많이 사용하는 framework는 Torch이나 TensorFlow일것입니다. 저는 Torch을 조금 더 주력으로 사용합니다. (대부분 연구자 분들도 Torch를 사용하시겠죠) 하지만 Torch 모델은 Arm cpu를 사용하는 device에서는 최적화되지않아 inference 속도가 다소 느립니다. 그래서 Torch모델을 Arm cpu 연산에 최적화된 TFLite 모델로 변환시키는 방법에 대해 말씀드리도록 하겠습니다.

 

Torch모델을 TFLite로 변환시키는 과정은 다음과 같습니다.

1. Torch 모델 → ONNX 모델
2. ONNX 모델 → OpenVINO 모델
3. OpenVINO모델 → TFLite 모델 

Torch모델은 detection 모델 중 하나인 yolov5를 사용하겠습니다. 위의 과정을 거치면서 변환이 정상적으로 되었는 지 확인하기 위해서 Torch모델일때의 detection결과와 변환과정에서 생긴 모델들의 detection결과와 같은 지 확인하면서 진행하겠습니다.

 

실험에 사용한 환경은 다음과 같습니다.

• Enviornment
  • torch: 1.10.1+cu111
  • tensorflow: 2.4.1
  • onnx: 1.8.0
  • openvino-dev: 2022.1.0
  • openvino2tensorflow: 1.31.3

위의 사용된 환경꼭 맞춰주셔야 코드 정상실행됩니다. (특히 torch버전 맞춰주세요!) 사용된 모든 코드는 여기서 사용가능합니다.

1. Torch 모델 → ONNX 모델

Torch 모델을 OpenVINO 또는 TensorFlow(TFLite)모델로 변환하기 위해서는 ONNX(Open Neural Network Exchange)라는 중간자적 모델을 거치게 됩니다. 그래서 Torch 모델을 ONNX모델로 변환하기 위해서는 다음과 같은 코드를 사용합니다. 

 

# yolov5_convert.py
model = attempt_load(save_path, device=device, inplace=True, fuse=True)  # load FP32 model
nc, names, stride = model.nc, model.names, model.stride

gs = int(max(model.stride))  # grid size (max stride)
imgsz = [check_img_size(x, gs) for x in imgsz]  # verify img_size are gs-multiples
im = torch.zeros(bs, 3, *imgsz).to(device)  # image size(1,3,640,640) BCHW iDetection

# torch2onnx.py
def torch2onnx(model, im, save_path, train, dynamic):
    try:
        logger.info(f'ONNX: starting export with onnx {onnx.__version__}...')
        torch.onnx.export(model, im, save_path, verbose=False, opset_version=12,
                            training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
                            do_constant_folding=not train,
                            input_names=['images'],
                            output_names=['output'],
                            dynamic_axes={'images': {0: 'batch'},  # shape(1,3,640,640)
                                        'output': {
                                            0: 'batch',
                                            }  # shape(1,25200,85)
                                        }if dynamic else None)
		...

yolov5 공식 repo에서는 attempt_load 함수로 model을 load한다는 것만 알아두시고 직접 모델을 load하고 싶으시면 자신만의 방법 (e.g. torch.load)으로 model을 load하시면 됩니다. Torch모델을 ONNX모델로 변환하기 위해서 torch.onnx.export함수를 사용하시면 됩니다. 중요하게 볼 parameter는 아래와 같습니다. 

• model: Torch model
• im: 모델의 input image shape을 결정하는 dummy tensor
• do_cosntant_folding: True로 설정 할 경우 Conv-BN layers가 Conv하나로 folding(fusing)되어 graph 최적화함
• dynamic_axes: batch size에 대한 static한 설정을 할지 말지를 결정 (저는 dynamic변수에 False를 할당했습니다)

 

위의 함수를 통해 yolov5s Torch모델을 ONNX모델로 변환하는 CLI와 그에 대한 결과는 다음과 같습니다. 아래와 같이 ONNX 모델로 변환된것을 확인 가능하며 ONNX모델로 detection한 결과가 Torch모델로 detection한 결과와 같은 것을 확인가능합니다.

ONNX 변환 확인

 

2. ONNX 모델 → OpenVINO 모델

OpenVINO는 Intel에서 개발한 AI 모델 inference와 deploy에 특화된 toolkit입니다. 아래와 같이 OpenVINO형태로 변환된 AI모델은 Intel CPU 대상으로 최적화된 연산이 가능합니다. 또한 다른 framework의 모델로부터 변환이 자유롭게 "잘" 된다는 특징때문에 저는 ONNX 모델에서 OpenVINO를 거쳐 TensorFlow 모델로 변환을 할 것입니다.

 

출처: https://www.intel.com/content/www/us/en/developer/tools/openvino-toolkit/overview.html

ONNX모델에서 OpenVINO모델로 변환하는 주요 코드는 다음과 같습니다. 

def onnx2openvino(model, onnx_path, save_path, data_type):
    try:
        logger.info(f'OpenVINO: starting export with openvino {ie.__version__}...')
        f = str(save_path).replace('.pt', f'_openvino_model{os.sep}')

        cmd = f"mo --input_model {onnx_path} --output_dir {f} --data_type {'FP16' if data_type=='fp16' else 'FP32'}"
        subprocess.check_output(cmd.split())  # converting

	...

OpenVINO모델로 변환하기 위해서 CLI를 사용하게 됩니다. CLI는 python의 subprocess를 통해 실행하였고 mo라는 명령어를 통해서 ONNX모델에서 OpenVINO모델로 변환하였습니다. (mo는 openvino-dev PyPI가 설치되어 있어야 함!) 변환에 필요한 parameter는 다음과 같습니다. 

• --input_model: ONNX model이 저장된 path
• --output_dir: OpenVINO모델이 저장될 directory
  • 기본적으로 OpenVINO모델은 .xml파일과 .bin파일로 나누어 저장됨
  • .xml 파일은 모델의 meta data를 저장한 파일, .bin파일은 모델의 weight값을 저장한 파일
• --data_tpye: Converting시에 weight의 data type을 결정
  • 저는 Float32을 사용하였음.

위의 함수를 통해 yolov5s ONNX모델을 OpenVINO모델로 변환하는 CLI와 그에 대한 결과는 다음과 같습니다. 아래와 같이 OpenVINO 모델이 잘 생성된것을 확인 가능하며 OpenVINO모델로 detection한 결과가 Torch모델로 detection한 결과와 같은 것을 확인가능합니다.

OpenVINO 변환 확인

3. OpenVINO 모델 → TFLite 모델

이제 우리가 원하는 TFLite로 모델로 변환하는 과정만 남았습니다. OpenVINO모델을 TFLite모델로 변환하기위해 PINTO님의 openvino2tensorflow PyPI를 사용할 것입니다. 변화 코드는 다음과 같습니다. 

def openvino2tflite(model_name, openvino_path, save_dir):
    try:
        logger.info(f'TFLite: starting starting export with TF {tf.__version__}...')
        if not Path(openvino_path).is_file():
            openvino_path = next(Path(openvino_path).glob('*.xml'))  # get *.xml file from *_openvino_model dir

        saved_model_path = str(save_dir / 'model_float32.tflite')
        cmd = f'openvino2tensorflow --model_path {openvino_path} --model_output_path {save_dir}\
                --output_no_quant_float32_tflite  --weight_replacement_config ./data/convert/replace_{model_name}.json'
        
        subprocess.check_output(cmd.split())        
    ...

이번에도 CLI를 통해서 변환을 합니다. openvino2tensorflow 명령어를 사용하며 주요 parameter는 다음과 같습니다.

• --model_path: OpenVINO 모델(.xml, .bin)이 저장된 directory
• --model_output_path: TFLite모델이 저장될 directory
• --output_no_quant_float32_tflite: Float32 data type을 TFLite모델로 변환함을 명시
• --weight_replacement_config: OpenVINO 모델에서 TFLite변환시 수정해야하는 configuration
  • 주로 Transpose나 Reshape연산을 할 때 차원 axis가 안맞을 경우 수동으로 convert시에 axis값을 수정해줘야함
  • 이는 OpenVINO와 TensorFlow간의 모델선언의 차이로 생기는 문제로 수동으로 수정해서 해결해줘야함

--weight_replacement_config가 이해가 잘 안되실 테니 좀 더 설명드립니다. TFLite 모델로 정상적으로 변환하기 위해서는  transpose연산의 dimension axis를 기존 [0,1,3,4,2]에서 [0,3,1,2,4]로 바꿔줘야 합니다. 그렇게 하기위해서 --weight_replacement_config의 parameter값으로  ./data/convert/replace_yolov5s.json을 설정하는데 해당 파일안의 내용은 아래와 같이 변경하고 싶은 layer_id를 선택하고 "values"의 key값으로 변경하려 하는 dimension axis를 설정해주면 됩니다. (아래는 하나의 transpose operation에 대한 dimension axis를 변경하는 내용을 설명드린것이고 전체 파일 내용은 제 github의 ./data/convert/replace_yolov5s.json을 참고해주세요.)

--weight_replacement_config 인자 설명

 

위의 openvino2tflite 함수를 통해 yolov5s OpenVINO모델을 TFLite모델로 변환하는 CLI와 그에 대한 결과는 다음과 같습니다. 아래와 같이 TFLite 모델이 잘 생성된것을 확인 가능하며 TFLite모델로 detection한 결과가 Torch모델로 detection한 결과와 같은 것을 확인가능합니다.

TFLite 변환 확인

 

4. 각 모델 별 Inference 속도 비교

속도 비교에 사용한 desktop의 spec은 다음과 같습니다. inference 속도를 측정하는데 사용되는 이미지는 2장(위의 bus.jpg와 zidane.jpg)입니다. Inference속도는 CPU, Float32 data type기준으로 출력하였습니다.

• Device
  • CPU: Intel(R) Xeon(R) Silver 4210R CPU @ 2.40GHz (40core)
  • GPU: GeForce RTX 3080 (4개) 

inference 속도 비교

위의 그림의 inference 속도에 대한 평균을 내어 정리한 테이블은 아래와 같습니다. 확실히 inference에 사용한 device의 cpu가 intel이므로 OpenVINO모델에서 가장 빠른것을 알 수 있고 arm cpu에 최적화된 TFLite 모델은 inference가 느린 것을 확인가능합니다. 

Model type	Inference speed (s)
Torch	0.088s
ONNX	0.077s
OpenVINO	0.042s
TFLite	0.360s

 

이후에는 arm cpu에 대해서도 inference비교할 것이며 quantization사용 시에 따른 속도 비교또한 진행하도록 하겠습니다. 감사합니당~~