参考：https://github.com/tflearn/tflearn/issues/964
解决方法：
""" Tensorflow graph freezer Converts Tensorflow trained models in .pb Code adapted from: https://gist.github.com/morgangiraud/249505f540a5e53a48b0c1a869d370bf#file-medium-tffreeze-1-py """ import os, argparse os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' import tensorflow as tf from tensorflow.python.framework import graph_util def freeze_graph(model_folder,output_graph="frozen_model.pb"): # We retrieve our checkpoint fullpath try: checkpoint = tf.train.get_checkpoint_state(model_folder) input_checkpoint = checkpoint.model_checkpoint_path print("[INFO] input_checkpoint:", input_checkpoint) except: input_checkpoint = model_folder print("[INFO] Model folder", model_folder) # Before exporting our graph, we need to precise what is our output node # This is how TF decides what part of the Graph he has to keep and what part it can dump output_node_names = "FullyConnected/Softmax" # NOTE: Change here # We clear devices to allow TensorFlow to control on which device it will load operations clear_devices = True # We import the meta graph and retrieve a Saver saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=clear_devices) # We retrieve the protobuf graph definition graph = tf.get_default_graph() input_graph_def = graph.as_graph_def() # We start a session and restore the graph weights with tf.Session() as sess: saver.restore(sess, input_checkpoint) # We use a built-in TF helper to export variables to constants output_graph_def = graph_util.convert_variables_to_constants( sess, # The session is used to retrieve the weights input_graph_def, # The graph_def is used to retrieve the nodes output_node_names.split(",") # The output node names are used to select the usefull nodes ) # Finally we serialize and dump the output graph to the filesystem 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)) print("[INFO] output_graph:",output_graph) print("[INFO] all done") if __name__ == '__main__': parser = argparse.ArgumentParser(description="Tensorflow graph freezer Converts trained models to .pb file", prefix_chars='-') parser.add_argument("--mfolder", type=str, help="model folder to export") parser.add_argument("--ograph", type=str, help="output graph name", default="frozen_model.pb") args = parser.parse_args() print(args," ") freeze_graph(args.mfolder,args.ograph) # However, before doing model.save(...) on TFLearn i have to do # ************************************************************ # del tf.get_collection_ref(tf.GraphKeys.TRAIN_OPS)[:] # ************************************************************ """ Then I call this command python tf_freeze.py --mfolder=<path_to_tflearn_model> Note The <path_to_tflearn_model> must not have the ".data-00000-of-00001". The output_node_names variable may change depending on your architecture. The thing is that you must reference the layer that has the softmax activation function. """
注意：
1、需要在 tflearn的model.save 前：
del tf.get_collection_ref(tf.GraphKeys.TRAIN_OPS)[:]
作用：去除模型里训练OP。
参考：https://github.com/tflearn/tflearn/issues/605#issuecomment-298478314
2、如果是有batch normalzition，或者残差网络层，会出现：
Error when loading the frozen graph with tensorflow.contrib.layers.python.layers.batch_norm
ValueError: graph_def is invalid at node u'BatchNorm/cond/AssignMovingAvg/Switch': Input tensor 'BatchNorm/moving_mean:0' Cannot convert a tensor of type float32 to an input of type float32_ref
freeze_graph.py doesn't seem to store moving_mean and moving_variance properly
An ugly way to get it working:
manually replace the wrong node definitions in the frozen graph
RefSwitch --> Switch + add '/read' to the input names
AssignSub --> Sub + remove use_locking attributes
则需要在restore模型后加入：
# fix batch norm nodes
for node in gd.node:
if node.op == 'RefSwitch':
node.op = 'Switch'
for index in xrange(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']
参考：https://github.com/tensorflow/tensorflow/issues/3628
I met the same issue when I was trying to export graph and variables by saved_model module. And finally I found a walk around to fix this issue:
Remove the TRAIN_OPS collections from graph collection. e.g.:
with dnn.graph.as_default():
del tf.get_collection_ref(tf.GraphKeys.TRAIN_OPS)[:]
The dumped graph may not be available for training again (by tflearn), but should be able to perform prediction and evaluation. This is useful when serving model by another module or language (e.g. tensorflow serving or tensorflow go binding). I'll do more further tests about this.
If you wanna re-train the model, please use the builtin "save" method and re-construction the graph and load the saved data when re-training.
2、可能需要在代码修改这行，
output_node_names = "FullyConnected/Softmax" # NOTE: Change here
参考：https://gist.github.com/morgangiraud/249505f540a5e53a48b0c1a869d370bf#file-medium-tffreeze-1-py
@vparikh10 @ratfury @rakashi I faced the same situation just like you.
From what I understood, you may have to change this line according to your network definition.
In my case, instead of having output_node_names = "Accuracy/prediction", I have output_node_names = "FullyConnected_2/Softmax".
I made this change after reading this suggestion
对我自己而言，写成softmax或者Softmax都是不行的！然后我将所有的node names打印出来：
打印方法:
with tf.Session() as sess: model = get_cnn_model(max_len, volcab_size) model.fit(trainX, trainY, validation_set=(testX, testY), show_metric=True, batch_size=1000, n_epoch=1) init_op = tf.initialize_all_variables() sess.run(init_op) for v in sess.graph.get_operations(): print(v.name)
然后确保output_node_names在里面。
附：gist里的代码，将output node names转换为参数
import os, argparse import tensorflow as tf # The original freeze_graph function # from tensorflow.python.tools.freeze_graph import freeze_graph dir = os.path.dirname(os.path.realpath(__file__)) def freeze_graph(model_dir, output_node_names): """Extract the sub graph defined by the output nodes and convert all its variables into constant Args: model_dir: the root folder containing the checkpoint state file output_node_names: a string, containing all the output node's names, comma separated """ if not tf.gfile.Exists(model_dir): raise AssertionError( "Export directory doesn't exists. Please specify an export " "directory: %s" % model_dir) if not output_node_names: print("You need to supply the name of a node to --output_node_names.") return -1 # We retrieve our checkpoint fullpath checkpoint = tf.train.get_checkpoint_state(model_dir) input_checkpoint = checkpoint.model_checkpoint_path # We precise the file fullname of our freezed graph absolute_model_dir = "/".join(input_checkpoint.split('/')[:-1]) output_graph = absolute_model_dir + "/frozen_model.pb" # We clear devices to allow TensorFlow to control on which device it will load operations clear_devices = True # We start a session using a temporary fresh Graph with tf.Session(graph=tf.Graph()) as sess: # We import the meta graph in the current default Graph saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=clear_devices) # We restore the weights saver.restore(sess, input_checkpoint) # We use a built-in TF helper to export variables to constants output_graph_def = tf.graph_util.convert_variables_to_constants( sess, # The session is used to retrieve the weights tf.get_default_graph().as_graph_def(), # The graph_def is used to retrieve the nodes output_node_names.split(",") # The output node names are used to select the usefull nodes ) # Finally we serialize and dump the output graph to the filesystem 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)) return output_graph_def if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("--model_dir", type=str, default="", help="Model folder to export") parser.add_argument("--output_node_names", type=str, default="", help="The name of the output nodes, comma separated.") args = parser.parse_args() freeze_graph(args.model_dir, args.output_node_names)