將tensorflow的ckpt轉換為pb

from tensorflow.python import pywrap_tensorflow
import tensorflow as tf
from tensorflow.python.framework import graph_util
'''
將節點名字打印出來
'''
def getAllNodes(checkpoint_path):
    reader = pywrap_tensorflow.NewCheckpointReader(checkpoint_path)
    var_to_shape_map = reader.get_variable_to_shape_map()
    # Print tensor name and values
    for key in var_to_shape_map:
        print("tensor_name: ", key)
        #print(reader.get_tensor(key))




def freeze_graph(ckpt, output_graph):
    #輸出節點的名稱，最直觀的是從tensorboard里讀，一般就是最后輸出的節點，例如這里就是輸出accuracy的節點
    output_node_names = 'FrameAccuracy/Cast'
    
    # saver = tf.train.import_meta_graph(ckpt+'.meta', clear_devices=True)
    saver = tf.compat.v1.train.import_meta_graph(ckpt + '.meta', clear_devices=True)
    graph = tf.get_default_graph()
    input_graph_def = graph.as_graph_def()

    with tf.Session() as sess:
        saver.restore(sess, ckpt)
        output_graph_def = graph_util.convert_variables_to_constants(
            sess=sess,
            input_graph_def=input_graph_def,
            output_node_names=output_node_names.split(',')
        )
        with tf.gfile.GFile(output_graph, 'wb') as fw:
            fw.write(output_graph_def.SerializeToString())
        print('{} ops in the final graph.'.format(len(output_graph_def.node)))

'''
把pb文件的節點讀出來
'''
def print_tensors(pb_file):
    print('Model File: {}\n'.format(pb_file))
    # read pb into graph_def
    with tf.gfile.GFile(pb_file, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

    # import graph_def
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def)

    # print operations
    for op in graph.get_operations():
        print(op.name + '\t' + str(op.values()))

'''
從ckpt中讀取圖結構，輸出可以被tensorboard讀取的圖文件
'''
def showNetFromCkpt(path):
    from tensorflow.python.platform import gfile
    graph = tf.get_default_graph()
    graphdef = graph.as_graph_def()
    _ = tf.train.import_meta_graph(path)
    #tensorboard的圖文件輸出的位置
    #使用tensorboard --logdir=E:\\MachineLearningProjects\\ViolentDetection_JD\\savedModels\\graph 進入tensorboard
    summary_write = tf.summary.FileWriter("E:\\MachineLearningProjects\\ViolentDetection_JD\\savedModels", graph)
    summary_write.close()

if __name__ == '__main__':
    #注意這里的path必須是絕對路徑！！
    ckpt_path='E:\\MachineLearningProjects\\ViolentDetection_JD\\savedModels\\save_epoch_38\\ViolenceNet.ckpt'

    #讀取圖文件，讀完了就注釋了就行，把輸出節點寫到上面的freeze_graph函數里
    #showNetFromCkpt(ckpt_path+".meta")
    #getAllNodes(ckpt_path)

    #將ckpt轉換為pb,這里寫pb的路徑，也必須是絕對路徑
    output_graph_path='E:\\MachineLearningProjects\\ViolentDetection_JD\\savedModels\\ViolenceNet.pb'
    freeze_graph(ckpt_path,output_graph_path)

    #將Pb文件的節點打印出來，看看有沒有問題
    print_tensors(output_graph_path)

參考文獻:

1. https://blog.csdn.net/u014090429/article/details/93486721
2. https://blog.csdn.net/guyuealian/article/details/82218092