"""
Tensorflow graph freezer
Converts Tensorflow trained models in .pbCode 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_utildef freeze_graph(model_folder,output_graph="frozen_model.pb"):# We retrieve our checkpoint fullpathtry:checkpoint = tf.train.get_checkpoint_state(model_folder)input_checkpoint = checkpoint.model_checkpoint_pathprint("[INFO] input_checkpoint:", input_checkpoint)except:input_checkpoint = model_folderprint("[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 dumpoutput_node_names = "FullyConnected/Softmax" # NOTE: Change here# We clear devices to allow TensorFlow to control on which device it will load operationsclear_devices = True# We import the meta graph and retrieve a Saversaver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=clear_devices)# We retrieve the protobuf graph definitiongraph = 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 constantsoutput_graph_def = graph_util.convert_variables_to_constants(sess,                        # The session is used to retrieve the weightsinput_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 filesystemwith 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>NoteThe <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.
"""

# 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']

with dnn.graph.as_default():del tf.get_collection_ref(tf.GraphKeys.TRAIN_OPS)[:]

    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)

import os, argparseimport 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 fileoutput_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 fullpathcheckpoint = tf.train.get_checkpoint_state(model_dir)input_checkpoint = checkpoint.model_checkpoint_path# We precise the file fullname of our freezed graphabsolute_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 operationsclear_devices = True# We start a session using a temporary fresh Graphwith tf.Session(graph=tf.Graph()) as sess:# We import the meta graph in the current default Graphsaver = 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 constantsoutput_graph_def = tf.graph_util.convert_variables_to_constants(sess, # The session is used to retrieve the weightstf.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 filesystemwith 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_defif __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)