TF Subclassed Model to TFLite

안녕하세요.

오늘은 Tensorflow Subclassed 모델을
TFLite로 변환하는 방법에 대해 다루겠습니다.

서브클래싱 모델은 tf.keras.Model을 상속받아 정의 모델입니다.
call 함수를 직접 구현해야합니다.

이렇게 구현한 모델의 학습 결과를 저장했습니다.
저의 경우 크기를 줄이기 위해 h5에 가중치만 저장했습니다.
모델을 불러와서 가중치를 입히는 식....

서론이 길었던 이유는 블로그나 스택오버플로우에 있는 건
대부분 서브클래싱 모델에서 동작 안하기 때문입니다

###############
# h5 to TFLite
###############

import tensorflow as tf
import os
import os.path as osp

DATA_PATH = ''

path = f'Models/{DATA_PATH}'

save_dir = DATA_PATH + '_TFLite'
os.makedirs(save_dir, exist_ok=True)

for i in range(10):   # 10-Fold
  cur_path = osp.join(path, 'my_model_' + str(i) + '-fold.h5')
  
  model = Model(len(LABELS))
  model.compute_output_shape(input_shape=())  # build, load_weight 시 에러
  model.load_weights(cur_path)
  
  # TFLite로 변환 후 저장
  converter = tf.lite.TFLiteConverter.from_keras_model(model)
  converter.optimizations = [tf.lite.Optimize.DEFAULT]  # 최적화 옵션
  tflite_model = converter.convert()
  open(f"{save_dir}/{i}.tflite", "wb").write(tflite_model)

이렇게 변환된 모델 파일은 불러들여,
추론에 사용할 수 있겠다.

추론은 다음과 같은 코드로 사용할 수 있고,
메서드 이름만 다르고, 추론 절차는 동일함을 확인할 수 있다.

import tensorflow as tf

model_path = ''

tflite_path = f'{model_path}/.tflite'
interpreter = tf.lite.Interpreter(model_path=tflite_path)   # 모델 불러오기
interpreter.allocate_tensors()  # 텐서할당

in_idx = interpreter.get_input_details()[0]['index']    # 입력 차원
out_idx = interpreter.get_output_details()[0]['index']  # 출력 차원
                                              
for features, labels in batch_test:     # 배치의 테스트 데이터
    interpreter.set_tensor(in_idx, features)    # 데이터 입력
    interpreter.invoke()        # model.predict(feature)
    best_pred = interpreter.get_tensor(out_idx)     # 모델의 출력 가져오기

후기 !

내가 코드를 잘못 짠건지...
모델 크기는 줄어들었으나, 추론 속도는 오히려 줄어들었다...
뭐가 문젠지 알게되면 포스팅을 수정하도록 하겠다..!
혹시 아신다면 댓글로 알려주세요!