TensorFlow模型转换h5转pb

在TensorFlow模型训练阶段一般使用model.save()将模型保存为h5格式，但部署阶段经常需要将训练好的模型固化为pb格式。

 h5模型转pb模型实现脚本： 

import tensorflow as tf

 from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2



 def h5_to_pb(h5_save_path):

    model = tf.keras.models.load_model(h5_save_path, compile=False)

    model.summary()

    full_model = tf.function(lambda Input: model(Input))

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



    # Get frozen ConcreteFunction

    frozen_func = convert_variables_to_constants_v2(full_model)

    frozen_func.graph.as_graph_def()



    layers = [op.name for op in frozen_func.graph.get_operations()]

    print("-" * 50)

    print("Frozen model layers: ")

    for layer in layers:

        print(layer)



    tf.io.write_graph(graph_or_graph_def=frozen_func.graph,

                      logdir="./frozen_models3",

                      name="model.pb",

                      as_text=False)

h5_to_pb("weights/model.h5")

如果h5模型只有参数没有模型结构需要重新构建模型并加载参数，然后另存为带模型结构的h5文件，之后再将模型转换为pb格式。

以resnet50为例：

from tensorflow.keras.applications.resnet50 import ResNet50

def create_model():

    base_model=ResNet50(include_top=True, weights=None, classes=2)

    model = tf.keras.Model(inputs=base_model.input, outputs=base_model.output)

    return model



model=create_model()



# 编译模型

model.compile(optimizer='adam',loss='categorical_crossentropy',metrics=['accuracy'])



model.load_weights('model_weights.h5')

model.save("weights/model.h5")