Skip to content

Instantly share code, notes, and snippets.

@morgangiraud
Last active August 21, 2023 12:46
Show Gist options
  • Save morgangiraud/249505f540a5e53a48b0c1a869d370bf to your computer and use it in GitHub Desktop.
Save morgangiraud/249505f540a5e53a48b0c1a869d370bf to your computer and use it in GitHub Desktop.
medium-tffreeze-1
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)
@jiumem
Copy link

jiumem commented Nov 30, 2016

Great guide! But go further, how can I remove all useless variable or op for inference when saving the metadata? I tried

saver = tf.train.Saver(tf.get_collection(tf.GraphKeys.VARIABLES, scope='inference'))

Seems helpless

@vparikh10
Copy link

I am getting an error -

AssertionError: Accuracy/predictions is not in graph

@ratfury
Copy link

ratfury commented May 5, 2017

@vparikh10 any ideas how to solve the problem?

@rakashi
Copy link

rakashi commented Aug 11, 2017

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

@BBarbosa
Copy link

BBarbosa commented Sep 6, 2017

@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."

@BBarbosa
Copy link

BBarbosa commented Sep 6, 2017

@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".

softmax

I made this change after reading this suggestion

@zewenli98
Copy link

I don't know where can I input my model file...HELP

@Aashit-Sharma
Copy link

@Servon-Lee u gotta pass the model dir instead of the file

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment