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)