【Tensorflow】Tensorflow checkpoint导入导出pb文件

为了把能够把tensorflow的模型放到tensorflow serving上，需要把checkpoint转换为pb文件

 

导出pb模型

这里踩了不少坑，首先需要获取输入和输出的tensor，在训练/测试文件之外重新写脚本来导出的话，需要重新定义输入输出对应的tensor，这样会比较麻烦，建议直接在test后面进行转换。代码如下：

prediction_signature_dict = tf.saved_model.signature_def_utils.build_signature_def(

    inputs={'input': tf.saved_model.utils.build_tensor_info(input_tensor)},
    outputs={'output': tf.saved_model.utils.build_tensor_info(output_tensor)},
    method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME,
)
loaded_graph = tf.Graph()
with tf.Session(graph=loaded_graph) as sess:
    # sess.run(init_op)
    loader = tf.train.import_meta_graph(checkpoint_path + '/model-x.meta')
    loader.restore(sess, checkpoint_path + 'model-x')
    builder = tf.saved_model.builder.SavedModelBuilder('export/')
    builder.add_meta_graph_and_variables(sess,
                                         [tf.saved_model.tag_constants.TRAINING],
                                         {'signature': prediction_signature_dict})
    builder.save()

其中，prediction_signature_dict定义是为了在读取文件时不需要使用原模型中输入输出tensor的name。

input_tensor和output_tensor为输入输出的tensor，即定义好之后需要被feed的tensor。

method_name对应的模型的类型

定义好之后按照代码中的流程即可导出

 

导入pb模型

导出后可以测试一下是否可用，代码如下：

with tf.Session(graph=tf.Graph()) as sess:
    meta_graph_def = tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.TRAINING], model_path)

    signature = meta_graph_def.signature_def

    x_tensor_name = signature['signature'].inputs['input'].name
    y_tensor_name = signature['signature'].outputs['output'].name

    input_x = sess.graph.get_tensor_by_name(x_tensor_name)
    output_y = sess.graph.get_tensor_by_name(y_tensor_name)

如果报错为graph中没有xxxtensor，说明你打包pb文件时给的input_tensor或者output_tensor有问题

在测试模型之前，记得将预处理部分自己手动加进去。

尝试了一下pix2pix-tensorflow的checkpoint转pb文件，最后出来的结果是的RGB排列的，要使用opencv时记得转成BGR。