在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")