Tensorflow中ckpt转pb模型代码

卷发擦边

于 2023-01-30 17:24:22 发布

阅读量433

点赞数

CC 4.0 BY-SA版权

分类专栏：深度学习模型部署文章标签： tensorflow

本文链接：https://blog.csdn.net/ELLIOTT_ALEX/article/details/128806085

深度学习模型部署专栏收录该内容

1 篇文章

订阅专栏

本文介绍如何将训练好的深度学习模型从ckpt格式转换为pb格式，以便于使用TF Serving框架进行部署。通过具体代码示例展示了整个转换过程。

工业界中，常常使用TFserving框架部署深度学习模型，为了能够将训练好的深度学习模型进行部署，通常需要将ckpt格式文件转为pb格式，以下为具体代码：

import os
import tensorflow as tf

checkpoint_file = tf.train.latest_checkpoint("./saved_mode/model_name")
graph = tf.Graph()

with graph.as_default():
    session_conf = tf.ConfigProto(allow_soft_placement=True,log_device_placement=False)
    sess = tf.Session(config=session_conf)
    with sess.as_default():
        saver = tf.train.import_meta_graph("{}.meta".format(checkpoint_file))
        saver.restore(sess,checkpoint_file)
        builder = tf.saved_model.builder.SavedModelBuilder("./saved_model/model_name/pb")
        input_tensor= tf.saved_model.utils.build_tensor_info(graph.get_tensor_by_name("input_tensor:0"))
        output_tensor= tf.saved_model.utils.build_tensor_info(graph.get_tensor_by_name("output_tensor:0"))
        labeling_signature = (tf.saved_model.signature_def_utils.build_signature_def(inputs={"input_tensor:0":input_tensor},outputs={"output_tensor:0":output_tensor},method_name="tensorflow/serving/predict"))
        builder.add_meta_grapg_and_variables(sess,[tf.saved_model.tag_constants.SERVING],signature_def_map={tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:labeling_signature})
        builder.save()
        print("Build complete!")