tensorflow2.x保存pb模型，用于opencv3.4.16 dnn模块调用

这个不开车的老司机

已于 2023-07-11 18:24:52 修改

阅读量3.6k

点赞数

分类专栏：模式识别 python 文章标签： dnn python 深度学习

于 2022-01-18 15:42:15 首次发布

本文链接：https://blog.csdn.net/non_hercules/article/details/122561465

版权

模式识别同时被 2 个专栏收录

7 篇文章 0 订阅

订阅专栏

python

2 篇文章 0 订阅

订阅专栏



# -------------------

from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2


def save_model_to_cv_dnn(netmodel, frozen_out_path, frozen_graph_filename):

    full_model = tf.function(lambda x: netmodel(x)).get_concrete_function(tf.TensorSpec(netmodel.inputs[0].shape, netmodel.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("-" * 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)  # 模型输出

    # 存储PB模型
    # Save frozen graph to disk
    tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                      logdir=frozen_out_path,
                      name=f"{frozen_graph_filename}.pb",
                      as_text=False)
    # Save its text representation
    tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                      logdir=frozen_out_path,
                      name=f"{frozen_graph_filename}.pbtxt",
                      as_text=True)

# path of the directory where you want to save your model
model.save('./data/tf_model_savedmodel', save_format="tf") print('export saved model.') model_loaded = tf.keras.models.load_model('./data/tf_model_savedmodel') yout = model_loaded.predict(x_test) ylab = yout[:, 0] > 0.5 print(ylab.shape) print(y_test.shape) model_loaded.evaluate(x = x_test,y = y_test)
frozen_out_path = './'  # 存储模型的路径
# name of the .pb file
frozen_graph_filename = "frozen_graph1"  # 模型名称
save_model_to_cv_dnn(model,frozen_out_path,frozen_graph_filename)