-
-
Save morgangiraud/249505f540a5e53a48b0c1a869d370bf to your computer and use it in GitHub Desktop.
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) |
I am getting an error -
AssertionError: Accuracy/predictions is not in graph
@vparikh10 any ideas how to solve the problem?
while running the code getting an error "AssertionError: Accuracy/predictions is not in graph"
Traceback (most recent call last):
File "createPB.py", line 55, in
freeze_graph(args.model_folder)
File "createPB.py", line 41, in freeze_graph
output_node_names.split(",") # The output node names are used to select the usefull nodes
File "/home/rakashi/tensorflow/local/lib/python2.7/site-packages/tensorflow/python/framework/graph_util_impl.py", line 202, in convert_variables_to_constants
inference_graph = extract_sub_graph(input_graph_def, output_node_names)
File "/home/rakashi/tensorflow/local/lib/python2.7/site-packages/tensorflow/python/framework/graph_util_impl.py", line 141, in extract_sub_graph
assert d in name_to_node_map, "%s is not in graph" % d
AssertionError: Accuracy/predictions is not in graph
@jiumem Try following this suggestion. I added it right before saving the trained model.
It solved me an error of "KeyError: "The name 'Adam' refers to an Operation not in the graph."
@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
I don't know where can I input my model file...HELP
@Servon-Lee u gotta pass the model dir instead of the file
Great guide! But go further, how can I remove all useless variable or op for inference when saving the metadata? I tried
Seems helpless