tensorflow模型ckpt转pb和测试推理

fzu_zhengbinchen

于 2020-10-12 21:57:30 发布

阅读量1.1k

点赞数

分类专栏： python 文章标签：深度学习 tensorflow

本文链接：https://blog.youkuaiyun.com/qq_26894673/article/details/109037916

版权

本文介绍如何将TensorFlow的checkpoint文件转换为protobuf（pb）模型，并详细阐述了使用pb文件进行模型推理的步骤。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

import cv2
import tensorflow as tf
from tensorflow.python.framework import graph_util


def freeze_graph(input_checkpoint ,output_graph):
    '''
    :param input_checkpoint:
    :param output_graph: PB模型保存路径
    :return:
    '''

    with tf.name_scope('input'):
        input_data = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_data')    #定义输入节点
    net = lanenet.LaneNet(phase='test', net_flag='vgg')
    ret = net.inference(input_tensor=input_data, name='lanenet_model')    # 网络模型结构
    output = tf.identity(ret, name='output_label')     # 定义输出节点

    saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
    graph = tf.get_default_graph() # 获得默认的图
    input_graph_def = graph.as_graph_def()  # 返回一个序列化的图代表当前的图

    with tf.Session() as sess:
        saver.restore(sess, input_checkp