[안드로이드/TF Lite] input/output shape 확인 법

안드로이드에서 텐서플로 라이트로 딥러닝 모델을 올릴 때

입/출력 shape 확인법

 

텐서플로 라이트 모델을 사용하기 전에 해당 모델의 입/출력 형태가 어떠한지 확인할 필요가 있다.

 

1. 먼저 모델의 경로를 명시하고 TFLite interpreter에 할당해준다.

# tflite model path
TFLITE_PATH = './TFLite/best_res18.tflite'

# tflite모델 로딩 및 텐서 할당
interpreter = tf.lite.Interpreter(model_path=TFLITE_PATH)

interpreter.allocate_tensors()

 

2. 입출력 텐서를 가져와서 shape을 print한다.

# 입출력 텐서 가져오기
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

input_shape = input_details[0]['shape']
output_shape = output_details[0]['shape']

print(input_shape.shape)
print(output_shape.shape)

그럼 해당 모델의 input shape과 output shape이 출력된다.

아래의 결과는 필자가 갖고 있는 모델의 입/출력 shape 이다.

>> [  1   3 256 256]
>> [ 1 20 64 64]

 

3. TFLite 모델의 결과까지 알기 위해서는 아래 코드를 작성한다.

# tflite output
output_data = interpreter.get_tensor(output_details[0]['index'])

output_data 변수의 모델의 output이 저장되었을 것이다.