Tensorflow训练的ckpt模型转pb代码

功能：将tensorflow1.几训练出来的ckpt模型转成pb模型

代码：

# -*- coding: utf-8 -*-
"""
@author: fancp
"""
import tensorflow.compat.v1 as tf
from nets import nets_factory
from tensorflow.python.framework import graph_util

FLAGS = tf.app.flags.FLAGS
tf.app.flags.DEFINE_string(
    'model_name', 'mobilenet_v2', 'The name of the architecture to train.')
tf.app.flags.DEFINE_integer('classNumber', 12, 'Dimensionality of the class.')
 
def freeze_graph(input_checkpoint,output_graph):
    network_fn = nets_factory.get_network_fn(
        FLAGS.model_name,
        num_classes= FLAGS.classNumber,
        is_training=False)
    input_shape = [1, 224, 224, 3]
    placeholder = tf.placeholder(name='input', dtype=tf.float32, shape=input_shape)
    logits, _ = network_fn(placeholder)
    # 指定输出的节点名称,该节点名称必须是原模型中存在的节点
    output_node_names = "MobilenetV2/Logits/output"
    sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
    saver = tf.train.Saver()
    saver.restore(sess, input_checkpoint) #恢复图并得到数据
    
    
    output_graph_def = graph_util.convert_variables_to_constants(  # 模型持久化，将变量值固定
        sess=sess,
        input_graph_def=sess.graph_def,
        output_node_names=output_node_names.split(","))# 如果有多个输出节点，以逗号隔开
 
    with tf.gfile.GFile(output_graph, "wb") as f: #保存模型
        f.write(output_graph_def.SerializeToString()) #序列化输出
    print("%d ops in the final graph." % len(output_graph_def.node)) #得到当前图有几个操作节点
 
 
if __name__ == '__main__':
    # 输入ckpt模型路径
    input_checkpoint='softmax/Top_MV2_soft.ckpt-248000'
    # 输出pb模型的路径
    out_pb_path="./frozen_model.pb"
    # 调用freeze_graph将ckpt转为pb
    freeze_graph(input_checkpoint,out_pb_path)