TensorFlow ckpt模型文件转换为pb模型文件

问题

ckpt转pb主要有两点要注意

1.要知道模型输入输出的节点名称

使用tools里的freeze_graph来读取ckpt，展示所有节点名称，查找搜需要的

2.转化后Value Error问题

Value Error错误信息一般为
ValueError: Input 0 of node … was passed float from … incompatible with expected float_ref.
对其进行类型转换即可

1.获取输入输出节点名称

使用tools里的freeze_graph来读取ckpt

保存ckpt模型的文件夹下的三个文件名
epoch_50.ckpt.data-00000-of-00001
epoch_50.ckpt.index
epoch_50.ckpt.meta

获取所有节点名称

from tensorflow.python.tools import freeze_graph

def freeze_graph_name(input_checkpoint):
    '''
    :param input_checkpoint: 
    '''
    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_checkpoint)
        output_graph_def = graph_util.convert_variables_to_constants(  
            sess=sess,
            input_graph_def=input_graph_def,# 等于:sess.graph_def
            output_node_names=[var.name[:-2] for var in tf.global_variables()])
        # 查看所有节点
        for op in graph.get_operations():
            print(op.name, op.values())

if __name__ == '__main__':
	input_checkpoint = 'models/ckpt/epoch_50.ckpt'
	freeze_graph_name(input_checkpoint)

看输出结果找到你需要的网络输入输出节点名称就行了

2.转化为pb文件并解决Value Error问题

from tensorflow.python.tools import freeze_graph

def freeze_graph(input_checkpoint, output_graph):
    '''
    :param input_checkpoint:
    :param output_graph: PB模型保存路径
    '''
    # 输出节点名称
    output_node_names = "out/pred" 
    
    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_checkpoint)
        
        # fix batch norm nodes
        # 解决value error问题
        for node in input_graph_def.node:
          if node.op == 'RefSwitch':
            node.op = 'Switch'
            for index in range(len(node.input)):
              if 'moving_' in node.input[index]:
                node.input[index] = node.input[index] + '/read'
          elif node.op == 'AssignSub':
            node.op = 'Sub'
            if 'use_locking' in node.attr: del node.attr['use_locking']

        output_graph_def = graph_util.convert_variables_to_constants(  
            sess=sess,
            input_graph_def=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__':
	input_checkpoint = 'models/ckpt/epoch_50.ckpt'
	output_graph = 'freeze_model.pb'
	freeze_graph(input_checkpoint, output_graph)

然后就可以直接使用pb文件预测了
预测程序自行搜索，很多
但要注意固化后的pb文件节点名称前默认加了import
就是在.get_tensor_by_name时 "out/pred"变成"import/out/pred"

参考链接：
https://blog.youkuaiyun.com/guyuealian/article/details/82218092
https://www.jb51.net/article/142183.htm
https://blog.youkuaiyun.com/derteanoo/article/details/89362635
https://github.com/onnx/tensorflow-onnx/issues/77