python - 将 tensorflow 检查点保存到 .pb protobuf 文件

Question

我训练了一个pix2pix tensorflow 上的模型，模型以检查点的形式保存在以下文件中:

model-15000.meta, model-15000.index, model-15000.data-00000-of-00001, >graph.pbtxt, 检查点.

现在，我想将其转换为 protobuf 文件 (.pb) 以用于部署目的。我遇到了 freeze_graph.py脚本来这样做，但我遇到了参数之一的麻烦，它是 output_node_names。

我尝试了几个图层名称，但出现以下错误:

AssertionError: generator/decoder_2/batchnorm/scale/gradients is not in graph

不确定如何找到 output_node_names

Answer 1

尝试使用下面的代码将 meta 文件转换为 pb 文件:

import tensorflow as tf
#Step 1 
#import the model metagraph
saver = tf.train.import_meta_graph('./model.meta', clear_devices=True)
#make that as the default graph
graph = tf.get_default_graph()
input_graph_def = graph.as_graph_def()
sess = tf.Session()
#now restore the variables
saver.restore(sess, "./model")

#Step 2
# Find the output name
graph = tf.get_default_graph()
for op in graph.get_operations(): 
  print (op.name)

#Step 3
from tensorflow.python.platform import gfile
from tensorflow.python.framework import graph_util

output_node_names="predictions_mod/Sigmoid"
output_graph_def = graph_util.convert_variables_to_constants(
        sess, # The session
        input_graph_def, # input_graph_def is useful for retrieving the nodes 
        output_node_names.split(",")  )    

#Step 4
#output folder
output_fld ='./'
#output pb file name
output_model_file = 'model.pb'
from tensorflow.python.framework import graph_io
#write the graph
graph_io.write_graph(output_graph_def, output_fld, output_model_file, as_text=False)

希望这有效!!!