在TensorFlow模型訓練階段一般使用model.save()將模型保存為h5格式，但部署階段經常需要將訓練好的模型固化為pb格式。

?h5模型轉pb模型實現腳本：?

import tensorflow as tffrom tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2def 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 ConcreteFunctionfrozen_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 ResNet50def 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 modelmodel=create_model()# 編譯模型model.compile(optimizer='adam',loss='categorical_crossentropy',metrics=['accuracy'])model.load_weights('model_weights.h5')model.save("weights/model.h5")

總結

以上是生活随笔為你收集整理的TensorFlow模型转换h5转pb的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。