[AI 모델 경량화] Representative Dataset 이용한 UINT8 양자화

지금가지 학습 후 양자화와 사전 양화를 진행해보았는데, 학습 후 양자화 하는 방법 중 한가지 빼 놓은 것이 있습니다. 바로 UINT8로 INT8보다 훨씬더 타입을 축소하게 만드는 것인데 이를 위해서는 몇가지 절차가 필요한데요 오늘은 그 방법에 대해서 살펴보도록 하겠습니다.

소스코드

from __future__ import absolute_import, division, print_function
import tensorflow as tf
import cv2
import numpy as np
import pathlib
import glob, os

BATCH_SIZE = 1
src = "/home/hero/data/datasets/representative/"
im_path = glob.glob(os.path.join(src, '*.jpg'))

def rep_data_gen():
    a = []
    global im_path
    for item in im_path:
        file_name = item
        img = cv2.imread(file_name)
        img = cv2.resize(img, (224, 224))
        img = img / 255.0
        img = img.astype(np.float32)
        a.append(img)
    a = np.array(a)
    print("===========================================================================")
    print("DONE!!!!!!!!!!!!!!!!!!!!")
    print("===========================================================================")
    img = tf.data.Dataset.from_tensor_slices(a).batch(1)
    for i in img.take(BATCH_SIZE):
        print(i)
        yield [i]

model = tf.keras.models.load_model('./saved')
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]

converter.target_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8

converter.representative_dataset = rep_data_gen

tflite_quant_model = converter.convert()

open("uint8_model.tflite", "wb").write(tflite_quant_model)

print("Model successfully converted into tflite file.")

결과

위와 같이 tflite로 변환되었는데, 주의할 것은 학습데이터를 미리 배치해 두고 참조할 수 있도록 Representative Data Set을 준비해 놓으며 데이터 셋은 1장 이상만 되도 결과는 동일하다. 실제 Inference Time의 경우 속도가 20~30% 정도 상승하는 것을 볼 수 있습니다. 다만 이러한 경량화는 정확도를 유지하면서 무조건 줄여주는 것은 아닐 수 있기에 반드시 정확도에 대한 확인이 필요하며 정확도라는 것은 학습하는 데이터와 학습하는 방식에 따라서도 달라질 수 있습니다.