TensorFlow模型转换h5转pb

最新推荐文章于 2022-11-25 11:00:58 发布

原创

最新推荐文章于 2022-11-25 11:00:58 发布 · 3.2k 阅读

17 ·

CC 4.0 BY-SA版权

文章标签：

#tensorflow #python #深度学习

本文介绍如何将TensorFlow的H5格式模型转换为PB格式，适用于模型部署。包括使用Python脚本进行转换的具体步骤，并提供了针对仅包含参数的H5模型的解决方案。

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