python 解析pb文件_tensorflow, ckpt 转 pb 并进行预测

tensorflow模型打包成pb

多个pb模型合并成一个

1. 查看节点名字

对 tf 不是很熟悉，所以有时间节点名字不清楚，用此办法找节点名字

# 处在调试状态，图已经建立起来后
graph = tf.get_default_graph().as_graph_def()

with open('graph', 'w', encoding='utf-8') as fgraph:
    fgraph.write(str(g.node))
    
# 然后根据程序中的节点名字去保存的文件中搜索，在后面生成 pb 文件写结点名字的时候用  
# saver.save(sess, 'model.ckpt') 里面保存的是变量，并不是所有图的节点

2. 生成 pb 文件

# output_node_names 为输出节点的名字
def freeze_graph(path='model.ckpt', output='model.pb'):
    saver = tf.train.import_meta_graph(path+'.meta', clear_devices=True)
    graph = tf.get_default_graph()
    input_graph_def = graph.as_graph_def()

    with tf.Session() as sess:
        saver.restore(sess, path)
        output_graph_def = graph_util.convert_variables_to_constants(
                           sess=sess,
                           input_graph_def=input_graph_def,   # = sess.graph_def,
                           output_node_names=['output/scores'])

        with tf.gfile.GFile(output, 'wb') as fgraph:
            fgraph.write(output_graph_def.SerializeToString())
            
# 程序中的输出节点
# with name_scope('output'):
#     self.scores = tf.*(..., name='scores')      

3. inference

主要是利用 get_tensor_by_name 来实现(注：此处节点名字要加上 :0）

with tf.gfile.GFile('model.pb', 'rb') as fgraph:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(fgraph.read())

with tf.Graph().as_default() as graph:
    tf.import_graph_def(graph_def, name='')

    input_x = graph.get_tensor_by_name('input_x:0')
    pred = graph.get_tensor_by_name('output/scores:0')

    sess = tf.Session(graph=graph)
    scores = sess.run(pred, feed_dict={input_x: x})

4. 合并多个图

用于将多个图合并到一起，这样只用一个 sess 就可以得到多个 model 输出

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

output_nodes = ['outputs']

def load_pb(path='model.pb'):
    with tf.gfile.GFile(path, 'rb') as fgraph:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(fgraph.read())

        return graph_def

def combined_graph():
    with tf.Graph().as_default() as g_combine:
        with tf.Session(graph=g_combine) as sess:
            graph_a = load_pb('graph_a.pb')
            graph_b = load_pb('graph_b.pb')

            tf.import_graph_def(graph_a, name='')
            tf.import_graph_def(graph_b, name='')

            g_combine_def = graph_util.convert_variables_to_constants(
                           sess=sess,
                           input_graph_def=sess.graph_def,
                           output_node_names=output_nodes)
            tf.train.write_graph(g_combine_def, './', 'model_combine.pb', as_text=False)