参考：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\nConverts 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,"\n")
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
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)
