Keras to TFlite

오늘은 keras 모델을 TFlite 모델로 변환하는 방법을 소개합니다.

tflite와 관련된 세부 정보는 Tensorflow 사이트를 방문하시면 쉽게 접하실 수 있습니다.

https://www.tensorflow.org/lite?hl=ko 

TensorFlow Lite | 휴대기기 및 에지 기기용 ML

보통 TFlite 모델을 변환하여 휴대용 기기 및 Arm core CPU에서 사용하곤합니다.

 

우선 변환을 위하여 Keras 모델을 준비합니다. 모델을 준비하는 방법은 아래의 함수를 사용하시면 됩니다.

• load_weights()
• load_model()

저번에도 설명드렸지만 모델 구조에 대한 정보가 있다면 load_weight으로 불러오시고 모델에 대한 정보가 없다면 load_model을 사용하는 것이 적절합니다.

 

이렇게 모델이 준비된다면 이제 tensorflow version을 고려하셔야 합니다.

단순하게 format을 tflite로 변경하시려 한다면 tensorflow 1.x 버전만 사용하시면 tflite 모델로 변경이 가능합니다.

하지만 quantization (양자화)를 적용하여 tflite 모델로 변경하려면 tensorflow 2.x 버전을 사용하셔야합니다.

 

[Quantization 양자화]

tflite 모델은 float, int에 대해 flow를 갖고 있기 때문에 어떤 자료형을 사용하셔도 상관 없습니다.

하지만 int형으로 quantize한 모델은 메모리와 연산 속도면에서 이점을 갖고 있습니다.

 

이제 코드를 살펴보겠습니다.

model.load_weights('keras.h5') # keras 모델을 로드합니다.

converter = tf.lite.TFLiteConverter.from_session(keras.backend.get_session(), model.inputs, model.outputs)
converter.optimizations = [tf.lite.Optimize.DEFAULT]

# Quantization function
def representative_data_gen():
	image = glob.glob('imageData.jpg')
    random.shuffle(image)
    
    for image in tqdm.tqdm(image):
    	img = cv2.imread(image)
        img = cv2.resize(img, dsize=input_shape, interpolation=cv2.INTER_AREA)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # 제가 사용한 모델은 흑백 이미지를 사용했기 때문에 GRAY로 변환했습니다.
        img = np.expand_dims(img, -1)
        img = np.expand_dims(img, 0)
        img = img.astype(np.float32) / 255.0
        yield [img]
        
converter.representative_dataset = representative_data_gen
converter.inference_input_type = tf.float32 # or int32
converter.inference_output_type = tf.float32
converted_model = converter.covert()

open(model_path + 'model.tflite', 'wb').write(converted_model)

만약 Quantize를 진행하지 않는다면 representative_gen function을 사용하지 않으시면 됩니다.

그리고 Quantize function을 사용한다면 입력 이미지로는 학습한 dataset을 가능한 많이 넣어주시는 것을 추천합니다.

 

저는 보통 해당 tflite 모델을 한번 더 처리를 진행하여 binary file로 변환하고 neural network engine을 통해 SoC에서 Inference를 위해 사용하였습니다.

 

추가로 python 환경에서 tflite 모델을 사용하는 방법을 적어봅니다.

# Python tf.lite model load
# Load the TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path = 'tf.lite model path')
interpreter.allocate_tensor()

# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# Test the model on random input data
input_shape = input_details[0]['shape']
input_data = np.array(np.random.random_sample(input_shape), type=np.float32)
interpreter.set_tensor(input_details[0]['index'], input_data)

interpreter.invoke()

# The function 'get_tensor()' return a copy of the tensor data.
# Use 'tensor()' in order to get a pointer to the tensor
output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data)

위의 Code를 통해 python 환경에서 tflite 모델을 사용할 수 있습니다.

 

저도 tflite 모델은 이번에 처음 변환을 하였으며 quantization 개념 또한 처음 공부하게되었습니다.

부족한 부분이 많지만 더 공부하고 정리되는 대로 추가로 수정하겠습니다.

 

감사합니다.