텐서플로우 모델을 frozen하고 optimized graph로 변환하기

tf_saved_model(.pb) 을 frozen fraph model(.pb , .pbtxt)로 변환하는 과정입니다.  tensorflow 모델을 tensorflowjs 모델로 변환하기 위해 필요한 중간과정으로 진행합니다.

1. Package 로드

import tensorflow as tf
from tensorflow import keras
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
import numpy as np

2. input 및 output 모델

Frozen Model

frozen_output_path = '' # frozen 모델을 저장할 경로
frozen_model = '' # frozen 모델 이름

tf_saved_model

model = tf.keras.models.load_model('') # tf_saved_model load

3. Keras 모델을 ConcreteFunction format으로 변환

full_model = tf.function(lambda x: model(x))
full_model = full_model.get_concrete_function(tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype)) 

4. 변수를 상수로 변환

frozen_func = convert_variables_to_constants_v2(full_model)
frozen_func.graph.as_graph_def()

5. 레이어 구조 출력

layers = [op.name for op in frozen_func.graph.get_operations()]
print("-" * 60)
print("Frozen model layers: ")
for layer in layers:
    print(layer)
print("-" * 60)
print("Frozen model inputs: ")
print(frozen_func.inputs)
print("Frozen model outputs: ")
print(frozen_func.outputs)

6. Frozen model 저장

tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                  logdir=frozen_output_path,
                  name=f'{frozen_model}.pb',
                  as_text=False)

tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                  logdir=frozen_output_path
                  name=f'{frozen_model}.pbtxt',
                  as_text=True)

 

 

 

출처

medium.com/@sebastingarcaacosta/how-to-export-a-tensorflow-2-x-keras-model-to-a-frozen-and-optimized-graph-39740846d9eb

 

How to export a TensorFlow 2.x Keras model to a frozen and optimized graph