def freeze_model():
model_path_suffix = os.path.join(FLAGS.network_def,
'input_{}_output_{}'.format(FLAGS.input_size, FLAGS.heatmap_size),
'joints_{}'.format(FLAGS.num_of_joints),
'stages_{}'.format(FLAGS.cpm_stages),
'init_{}_rate_{}_step_{}'.format(FLAGS.init_lr, FLAGS.lr_decay_rate,
FLAGS.lr_decay_step)
)
model_save_dir = os.path.join('models',
'weights',
model_path_suffix)
model_path = os.path.join(model_save_dir, FLAGS.model_path)
model_path = 'models/weights/cpm_hand'
# Load graph and dump to protobuf
meta_graph = tf.train.import_meta_graph(model_path + '.meta')
tf.train.write_graph(tf.get_default_graph(), 'frozen_models/', 'graph_proto.pb')
output_graph_path = os.path.join('frozen_models', '{}_frozen.pb'.format('cpm_hand'))
freeze_graph(input_graph='frozen_models/graph_proto.pb',
input_saver='',
input_checkpoint=model_path,
output_graph=output_graph_path,
output_node_names=FLAGS.output_node_names,
restore_op_name='save/restore_all',
clear_devices=True,
initializer_nodes='',
variable_names_blacklist='',
input_binary=False,
filename_tensor_name='save/Const:0')
def testFreezeGraph(self):
checkpoint_prefix = os.path.join(self.get_temp_dir(), "saved_checkpoint")
checkpoint_state_name = "checkpoint_state"
input_graph_name = "input_graph.pb"
output_graph_name = "output_graph.pb"
# We'll create an input graph that has a single variable containing 1.0,
# and that then multiplies it by 2.
with tf.Graph().as_default():
variable_node = tf.Variable(1.0, name="variable_node")
output_node = tf.mul(variable_node, 2.0, name="output_node")
sess = tf.Session()
init = tf.initialize_all_variables()
sess.run(init)
output = sess.run(output_node)
self.assertNear(2.0, output, 0.00001)
saver = tf.train.Saver()
saver.save(sess, checkpoint_prefix, global_step=0,
latest_filename=checkpoint_state_name)
tf.train.write_graph(sess.graph.as_graph_def(), self.get_temp_dir(),
input_graph_name)
# We save out the graph to disk, and then call the const conversion
# routine.
input_graph_path = os.path.join(self.get_temp_dir(), input_graph_name)
input_saver_def_path = ""
input_binary = False
input_checkpoint_path = checkpoint_prefix + "-0"
output_node_names = "output_node"
restore_op_name = "save/restore_all"
filename_tensor_name = "save/Const:0"
output_graph_path = os.path.join(self.get_temp_dir(), output_graph_name)
clear_devices = False
freeze_graph.freeze_graph(input_graph_path, input_saver_def_path,
input_binary, input_checkpoint_path,
output_node_names, restore_op_name,
filename_tensor_name, output_graph_path,
clear_devices, "")
# Now we make sure the variable is now a constant, and that the graph still
# produces the expected result.
with tf.Graph().as_default():
output_graph_def = tf.GraphDef()
with open(output_graph_path, "rb") as f:
output_graph_def.ParseFromString(f.read())
_ = tf.import_graph_def(output_graph_def, name="")
self.assertEqual(4, len(output_graph_def.node))
for node in output_graph_def.node:
self.assertNotEqual("Variable", node.op)
with tf.Session() as sess:
output_node = sess.graph.get_tensor_by_name("output_node:0")
output = sess.run(output_node)
self.assertNear(2.0, output, 0.00001)
def convert_tf_session_bundle(session_bundle_dir,
output_node_names,
output_dir,
quantization_dtype=None,
skip_op_check=False,
strip_debug_ops=False):
"""Freeze the Session Bundle model and check the model compatibility with
Tensorflow.js.
Optimize and convert the model to Tensorflow.js format, when the model passes
the compatiblity check.
Args:
session_bundle_dir: string The session bundle model directory.
output_node_names: string The names of the output nodes, comma separated.
output_dir: string The name of the output directory. The directory
will consist of
- a file named 'tensorflowjs_model.pb'
- a JSON weights manifest file named 'weights_manifest.json'
- possibly sharded binary weight files.
quantization_dtype: An optional numpy dtype to quantize weights to for
compression. Only np.uint8 and np.uint16 are supported.
skip_op_check: Bool whether to skip the op check.
strip_debug_ops: Bool whether to strip debug ops.
"""
print("Tensorflow has deprecated the Session Bundle format, ",
"please migrate to SavedModel.")
if not os.path.exists(output_dir):
os.makedirs(output_dir)
output_graph = os.path.join(output_dir, DEFAULT_MODEL_PB_FILENAME)
checkpoint = tf.train.get_checkpoint_state(session_bundle_dir)
input_checkpoint = checkpoint.model_checkpoint_path
frozen_file = output_graph + '.frozen'
freeze_graph.freeze_graph(
'',
'',
True,
input_checkpoint,
output_node_names,
'',
'',
frozen_file,
True,
'',
input_meta_graph=input_checkpoint + '.meta')
graph = load_graph(output_graph + '.frozen', output_node_names)
try:
optimize_graph(graph, output_graph, quantization_dtype=quantization_dtype,
skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops)
finally:
# Clean up the temp files.
if os.path.exists(frozen_file):
os.remove(frozen_file)
def convert_tf_saved_model(saved_model_dir, output_node_names,
output_dir, saved_model_tags='serve',
quantization_dtype=None,
skip_op_check=False,
strip_debug_ops=False):
"""Freeze the SavedModel and check the model compatibility with Tensorflow.js.
Optimize and convert the model to Tensorflow.js format, when the model passes
the compatiblity check.
Args:
saved_model_dir: string The saved model directory.
output_node_names: string The names of the output nodes, comma separated.
output_dir: string The name of the output directory. The directory
will consist of
- a file named 'tensorflowjs_model.pb'
- a JSON weights manifest file named 'weights_manifest.json'
- possibly sharded binary weight files.
saved_model_tags: string Tagset of the MetaGraphDef to load, in comma
separated string format. Defaulted to 'serve'
quantization_dtype: An optional numpy dtype to quantize weights to for
compression. Only np.uint8 and np.uint16 are supported.
skip_op_check: Bool whether to skip the op check.
strip_debug_ops: Bool whether to strip debug ops.
"""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
output_graph = os.path.join(output_dir, DEFAULT_MODEL_PB_FILENAME)
frozen_file = output_graph + '.frozen'
freeze_graph.freeze_graph(
'',
'',
True,
'',
output_node_names,
'',
'',
frozen_file,
True,
'',
saved_model_tags=saved_model_tags,
input_saved_model_dir=saved_model_dir)
graph = load_graph(output_graph + '.frozen', output_node_names)
try:
optimize_graph(graph, output_graph, quantization_dtype=quantization_dtype,
skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops)
finally:
# Clean up the temp files.
if os.path.exists(frozen_file):
os.remove(frozen_file)
def _export_graph(self):
"""
Exports latest saved model to .bytes format for Unity embedding.
"""
target_nodes = ','.join(self._process_graph())
ckpt = tf.train.get_checkpoint_state(self.model_path)
freeze_graph.freeze_graph(input_graph=self.model_path + '/raw_graph_def.pb',
input_binary=True,
input_checkpoint=ckpt.model_checkpoint_path,
output_node_names=target_nodes,
output_graph=self.model_path + '/' + self.env_name + "_" + self.run_id + '.bytes',
clear_devices=True, initializer_nodes="", input_saver="",
restore_op_name="save/restore_all", filename_tensor_name="save/Const:0")
def testFreezeMetaGraph(self):
tmp_dir = self.get_temp_dir()
checkpoint_prefix = os.path.join(tmp_dir, "meta_graph_checkpoint")
checkpoint_state_name = "checkpoint_state"
output_graph_filename = os.path.join(tmp_dir, "output_graph.pb")
with ops.Graph().as_default():
variable_node = variables.VariableV1(1.0, name="variable_node")
output_node = math_ops.multiply(variable_node, 2.0, name="output_node")
sess = session.Session()
init = variables.global_variables_initializer()
sess.run(init)
output = sess.run(output_node)
self.assertNear(2.0, output, 0.00001)
saver = saver_lib.Saver()
checkpoint_path = saver.save(
sess,
checkpoint_prefix,
global_step=0,
latest_filename=checkpoint_state_name)
input_saver_def_path = ""
input_binary = True
output_node_names = "output_node"
restore_op_name = "save/restore_all"
filename_tensor_name = "save/Const:0"
clear_devices = False
input_meta_graph = checkpoint_path + ".meta"
freeze_graph.freeze_graph(
"", input_saver_def_path, input_binary, checkpoint_path,
output_node_names, restore_op_name, filename_tensor_name,
output_graph_filename, clear_devices, "", "", "", input_meta_graph)
# Now we make sure the variable is now a constant, and that the graph still
# produces the expected result.
with ops.Graph().as_default():
output_graph_def = graph_pb2.GraphDef()
with open(output_graph_filename, "rb") as f:
output_graph_def.ParseFromString(f.read())
_ = importer.import_graph_def(output_graph_def, name="")
self.assertEqual(4, len(output_graph_def.node))
for node in output_graph_def.node:
self.assertNotEqual("VariableV2", node.op)
self.assertNotEqual("Variable", node.op)
with session.Session() as sess:
output_node = sess.graph.get_tensor_by_name("output_node:0")
output = sess.run(output_node)
self.assertNear(2.0, output, 0.00001)
def export_graph(model_path, env_name="env", target_nodes="action,value_estimate,action_probs"):
"""
Exports latest saved model to .bytes format for Unity embedding.
:param model_path: path of model checkpoints.
:param env_name: Name of associated Learning Environment.
:param target_nodes: Comma separated string of needed output nodes for embedded graph.
"""
ckpt = tf.train.get_checkpoint_state(model_path)
freeze_graph.freeze_graph(input_graph=model_path + '/raw_graph_def.pb',
input_binary=True,
input_checkpoint=ckpt.model_checkpoint_path,
output_node_names=target_nodes,
output_graph=model_path + '/' + env_name + '.bytes',
clear_devices=True, initializer_nodes="", input_saver="",
restore_op_name="save/restore_all", filename_tensor_name="save/Const:0")
def export_graphdef(filename):
from tensorflow.python.tools import freeze_graph
graph = tf.get_default_session().graph
sess = tf.get_default_session()
graph_def = graph.as_graph_def()
# fix batch norm nodes
for node in graph_def.node:
if node.op == 'RefSwitch':
node.op = 'Switch'
for index in range(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']
elif node.op == 'AssignAdd':
node.op = 'Add'
if 'use_locking' in node.attr: del node.attr['use_locking']
inputs = []
for inp in node.input:
if inp[0] == '^':
node.input.remove(inp)
saver_path = tf.train.Saver().save(sess, 'cache/karras2019stylegan-ffhq-1024x1024.ckpt')
converted_graph = tf.graph_util.convert_variables_to_constants(sess, graph_def, ['Gs/images_out'])
tf.train.write_graph(converted_graph, 'cache', f'{filename}_converted.pb', as_text=False)
graph_path = tf.train.write_graph(converted_graph, 'cache', f'{filename}.pbtxt')
print('Freezing graph')
freeze_graph.freeze_graph(
input_graph=graph_path,
input_saver='',
input_binary=False,
input_checkpoint=saver_path,
output_node_names=['Gs/images_out'],
restore_op_name='',
filename_tensor_name='',
output_graph=f'cache/frozen_{filename}.pb',
clear_devices=False,
initializer_nodes='',
variable_names_whitelist="",
variable_names_blacklist="",
input_meta_graph=None,
input_saved_model_dir=None
)
def tf_to_pb(sess, checkpoint, output, output_dir=None):
"""
Saves a frozen tensorflow graph (a protobuf file).
See also https://leimao.github.io/blog/Save-Load-Inference-From-TF-Frozen-Graph/
Parameters
----------
sess : tensorflow session
session with graph to be saved
checkpoint : string
checkpoint of tensorflow model to be converted to protobuf (output will be <checkpoint>.pb)
output : list of strings
list of the names of output nodes (is returned by load_models)
output_dir : string, optional
path to the directory that exported models should be saved to.
If None, will export to the directory of the checkpoint file.
"""
output_dir = (os.path.expanduser(output_dir)
if output_dir else os.path.dirname(checkpoint))
ckpt_base = os.path.basename(checkpoint)
# save graph to pbtxt file
pbtxt_file = os.path.normpath(output_dir + "/" + ckpt_base + ".pbtxt")
tf.train.write_graph(sess.graph.as_graph_def(),
"",
pbtxt_file,
as_text=True)
# create frozen graph from pbtxt file
pb_file = os.path.normpath(output_dir + "/" + ckpt_base + ".pb")
freeze_graph.freeze_graph(
input_graph=pbtxt_file,
input_saver="",
input_binary=False,
input_checkpoint=checkpoint,
output_node_names=",".join(output),
restore_op_name="save/restore_all",
filename_tensor_name="save/Const:0",
output_graph=pb_file,
clear_devices=True,
initializer_nodes="",
)
def convert_tf_session_bundle(session_bundle_dir, output_node_names,
output_dir):
"""Freeze the Session Bundle model and check the model compatibility with
Tensorflow.js.
Optimize and convert the model to Tensorflow.js format, when the model passes
the compatiblity check.
Args:
session_bundle_dir: string The session bundle model directory.
output_node_names: string The names of the output nodes, comma separated.
output_dir: string The name of the output directory. The directory
will consist of
- a file named 'tensorflowjs_model.pb'
- a JSON weights manifest file named 'weights_manifest.json'
- possibly sharded binary weight files.
"""
print("Tensorflow has deprecated the Session Bundle format, ",
"please migrate to SavedModel.")
if not os.path.exists(output_dir):
os.makedirs(output_dir)
output_graph = os.path.join(output_dir, DEFAULT_MODEL_PB_FILENAME)
checkpoint = tf.train.get_checkpoint_state(session_bundle_dir)
input_checkpoint = checkpoint.model_checkpoint_path
frozen_file = output_graph + '.frozen'
freeze_graph.freeze_graph('',
'',
True,
input_checkpoint,
output_node_names,
'',
'',
frozen_file,
True,
'',
input_meta_graph=input_checkpoint + '.meta')
graph = load_graph(output_graph + '.frozen', output_node_names)
unsupported = validate(graph.as_graph_def().node)
if unsupported:
print('Unsupported Ops in the model\n' + ', '.join(unsupported))
else:
optimize_graph(graph, output_graph)
# Clean up the temp files.
if os.path.exists(frozen_file):
os.remove(frozen_file)
def save_as_pb(self, directory, filename):
if not os.path.exists(directory):
os.makedirs(directory)
# Save check point for graph frozen later
ckpt_filepath = self.save(directory=directory, filename=filename)
pbtxt_filename = filename + ".pbtxt"
pbtxt_filepath = os.path.join(directory, pbtxt_filename)
pb_filepath = os.path.join(directory, filename + ".pb")
# This will only save the graph but the variables will not be saved.
# You have to freeze your model first.
tf.train.write_graph(
graph_or_graph_def=self.sess.graph_def,
logdir=directory,
name=pbtxt_filename,
as_text=True,
)
# Freeze graph
# Method 1
freeze_graph.freeze_graph(
input_graph=pbtxt_filepath,
input_saver="",
input_binary=False,
input_checkpoint=ckpt_filepath,
output_node_names="cnn/output",
restore_op_name="save/restore_all",
filename_tensor_name="save/Const:0",
output_graph=pb_filepath,
clear_devices=True,
initializer_nodes="",
)
# Method 2
"""
graph = tf.get_default_graph()
input_graph_def = graph.as_graph_def()
output_node_names = ['cnn/output']
output_graph_def = graph_util.convert_variables_to_constants(self.sess,
input_graph_def, output_node_names)
with tf.gfile.GFile(pb_filepath, 'wb') as f:
f.write(output_graph_def.SerializeToString())
"""
return pb_filepath
def convert_tf_saved_model(saved_model_dir,
output_node_names,
output_dir,
saved_model_tags='serve'):
"""Freeze the SavedModel and check the model compatibility with Tensorflow.js.
Optimize and convert the model to Tensorflow.js format, when the model passes
the compatiblity check.
Args:
saved_model_dir: string The saved model directory.
output_node_names: string The names of the output nodes, comma separated.
output_dir: string The name of the output directory. The directory
will consist of
- a file named 'tensorflowjs_model.pb'
- a JSON weights manifest file named 'weights_manifest.json'
- possibly sharded binary weight files.
saved_model_tags: string Tagset of the MetaGraphDef to load, in comma
separated string format. Defaulted to 'serve'
"""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
output_graph = os.path.join(output_dir, DEFAULT_MODEL_PB_FILENAME)
frozen_file = output_graph + '.frozen'
freeze_graph.freeze_graph('',
'',
True,
'',
output_node_names,
'',
'',
frozen_file,
True,
'',
saved_model_tags=saved_model_tags,
input_saved_model_dir=saved_model_dir)
graph = load_graph(output_graph + '.frozen', output_node_names)
unsupported = validate(graph.as_graph_def().node)
if unsupported:
print('Unsupported Ops in the model\n' + ', '.join(unsupported))
else:
optimize_graph(graph, output_graph)
# Clean up the temp files.
if os.path.exists(frozen_file):
os.remove(frozen_file)
def test_1(self):
with open(self.class_label_file, 'w+') as labels_file:
for a in range(10):
labels_file.write(str(a + 1) + "\n")
image_size = 224
with tf.Graph().as_default():
batch_size, height, width, channels = 1, image_size, image_size, 3
images = tf.random.uniform([batch_size, height, width, channels],
maxval=1)
# Create the model.
i_placeholder = tf.placeholder(
name='input',
dtype=tf.float32,
shape=[1, image_size, image_size, 3])
net = self.my_conv_2d(i_placeholder, [1, 3, 3, 1], 1, 1, 'first')
net = tf.nn.avg_pool2d(net,
224,
strides=1,
padding='VALID',
name='AvgPool_1a')
net = self.my_conv_2d(net, [1, 1, 1, 10],
10,
1,
'fc',
activation_fn=None,
with_bias_add=False)
net = tf.squeeze(net, [1, 2], name='SpatialSqueeze')
probabilities = tf.nn.softmax(net, name='Softmax')
saver = tf.train.Saver()
init_op = tf.global_variables_initializer()
with tf.Session() as sess:
sess.run(init_op)
probabilities = sess.run(probabilities,
{i_placeholder: images.eval()})
save_path = saver.save(sess, self.checkpoint_file)
with gfile.GFile(self.graph_file, 'wb') as f:
f.write(sess.graph_def.SerializeToString())
freeze_graph.freeze_graph(self.graph_file, '', True,
self.checkpoint_file, 'Softmax', '',
'', self.frozen_graph_file, False,
'')
def freeze_from_checkpoint(model_dir, checkpoint_subdir, model_name):
checkpoint_dir = os.path.join(model_dir, checkpoint_subdir)
checkpoint_path = tf.train.latest_checkpoint(checkpoint_dir)
input_graph_path = os.path.join(checkpoint_dir, 'graph.pbtxt')
input_saver_def_path = ""
input_binary = False
output_node_names = "masked_pred_test,masked_scale_test"
restore_op_name = "save/restore_all"
filename_tensor_name = "save/Const:0"
output_frozen_graph_name = os.path.join(model_dir, 'frozen_' + model_name + '.pb')
clear_devices = True
freeze_graph.freeze_graph(input_graph_path, input_saver_def_path,
input_binary, checkpoint_path, output_node_names,
restore_op_name, filename_tensor_name,
output_frozen_graph_name, clear_devices, "")
def freeze():
'''
cmd창에서도 할 수 있다.
> python freeze_graph.py
--input_graph = ./model_pb/my_graph.pbtxt --input_checkpoint = ./model_pb/model.ckpt --output_graph=./model_pb/my_graph.pb
'''
input_graph = './model_pb/my_graph.pbtxt'
input_checkpoint = './model_pb/model.ckpt'
output_graph = './model_pb/my_graph.pb' # 저장할 파일 이름
output_node_names = 'L2/Sigmoid' # 어디까지 내 보낼지
freeze_graph.freeze_graph(input_graph, "", False, input_checkpoint,
output_node_names, None, None, output_graph,
True, None)
def convert_keras_to_pb(keras_model, out_names, models_dir, model_filename):
model = load_model(keras_model)
out_names = model.outputs
print(out_names)
model.summary()
K.set_learning_phase(0)
sess = K.get_session()
saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V2)
checkpoint_path = saver.save(sess,
'saved_ckpt',
global_step=0,
latest_filename='checkpoint_state')
graph_io.write_graph(sess.graph, '.', 'tmp.pb')
freeze_graph.freeze_graph('./tmp.pb', '', False, checkpoint_path,
out_names, "save/restore_all", "save/Const:0",
models_dir + model_filename, False, "")
def export_model(input_node_names, output_node_name):
freeze_graph.freeze_graph('out/' + MODEL_NAME + '.pbtxt', None, False,
'out/' + MODEL_NAME + '.chkp', output_node_name,
"save/restore_all", "save/Const:0",
'out/frozen_' + MODEL_NAME + '.pb', True, "")
input_graph_def = tf.GraphDef()
with tf.gfile.Open('out/frozen_' + MODEL_NAME + '.pb', "rb") as f:
input_graph_def.ParseFromString(f.read())
output_graph_def = optimize_for_inference_lib.optimize_for_inference(
input_graph_def, [input_node_names], [output_node_name],
tf.float32.as_datatype_enum)
with tf.gfile.FastGFile('out/opt_' + MODEL_NAME + '.pb', "wb") as f:
f.write(output_graph_def.SerializeToString())
def freeze_keras_model(keras_model_path):
# Load keras model
model = load_model(keras_model_path)
# Observe the input_node_name and output_node_name which are used for creating inference graph with tensorrt
print(model.inputs)
print(model.outputs)
saver = tf.train.Saver(write_version=tf.train.SaverDef.V2)
with K.get_session() as sess:
# saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V2)
checkpoint_path = saver.save(sess, FLAGS.saved_ckpt,global_step=0,latest_filename='checkpoint_state')
graph_io.write_graph(sess.graph,'.',FLAGS.graphdef_file)
print(checkpoint_path)
freeze_graph.freeze_graph(FLAGS.graphdef_file, '',
False, checkpoint_path, FLAGS.output_node_name,
"save/restore_all", "save/Const:0",
FLAGS.frozen_graph, False, "")
def freeze_graph(sess, ckpt, output):
print("Loading checkpoint...")
saver = tf.train.Saver()
saver.restore(sess, ckpt)
print("Writing graph...")
if not os.path.isdir("_Cache"):
os.makedirs("_Cache")
_dir = os.path.join("_Cache", "Model")
saver.save(sess, _dir)
graph_io.write_graph(sess.graph, "_Cache", "Model.pb", False)
print("Freezing graph...")
freeze_graph.freeze_graph(os.path.join("_Cache", "Model.pb"), "", True,
os.path.join("_Cache", "Model"), output,
"save/restore_all", "save/Const:0",
"Frozen.pb", True, "")
print("Done")
def exportFrozenGraphForInference(self, clear_devices=True):
# Freeze the graph
input_graph_path = os.path.join(self.model_ckpt_directory,
self.graph_file_name)
ckpt_path = tf.train.latest_checkpoint(self.model_ckpt_directory)
input_saver_def_path = ""
output_node_name = "inference_model/" + self.y_inference
print('Freezing graph {}'.format(input_graph_path))
freeze_graph.freeze_graph(input_graph_path, input_saver_def_path,
self.use_binary_export, ckpt_path,
output_node_name, self.restore_op_name,
self.frozen_filename,
self.output_frozen_graph_name, clear_devices,
"")
def freeze(self):
init = tf.global_variables_initializer()
self._tensorflow_session.run(init)
self.checkpoint.load_all()
saver = tf.train.Saver()
saver.save(self._tensorflow_session, './gazeml.ckpt')
tf.train.write_graph(self._tensorflow_session.graph.as_graph_def(),
'.',
'gazeml.pbtxt',
as_text=True)
from tensorflow.python.tools import freeze_graph
freeze_graph.freeze_graph(
'./gazeml.pbtxt', '', False, './gazeml.ckpt',
'hourglass/hg_2/after/hmap/conv/BiasAdd:0,upscale/mul:0,radius/out/fc/BiasAdd:0',
'save/restore_all', 'save/Const:0', './gazeml.pb', True, '')
def freezer(sm, input_ckpt, output_graph):
output_node_names = 'train/c'
initializer_nodes = ''
#tags = tag_constants.SERVING,
tags = 'serve'
input_ckpt = tf.train.latest_checkpoint(ckpt_dir)
#meta_graph_file = input_ckpt + '.meta'
meta_graph_file = input_ckpt + '.meta'
print('meta_graph_file: ', meta_graph_file)
with tf.Session() as sess:
new_graph_location = freeze_graph.freeze_graph(
input_graph=utils.get_only_graph_def_from_sm(sm),
input_saver='',
input_binary=True,
input_checkpoint=input_ckpt,
output_node_names=output_node_names,
restore_op_name=None,
filename_tensor_name=None,
output_graph=output_graph,
clear_devices=True,
initializer_nodes=initializer_nodes,
#input_meta_graph = sm.meta_graphs[0],
#input_meta_graph = False,
input_meta_graph=meta_graph_file,
input_saved_model_dir=model_dir,
saved_model_tags=tags,
#checkpoint_version = saver_pb2.SaverDef.V2,
)
print('the new graph is at: ', new_graph_location)
print('graph freezed')
def export_model(self):
"""
Exports latest saved model to .tf format for Unity embedding.
"""
with self.graph.as_default():
target_nodes = ','.join(self._process_graph())
ckpt = tf.train.get_checkpoint_state(self.model_path)
freeze_graph.freeze_graph(
input_graph=self.model_path + '/raw_graph_def.pb',
input_binary=True,
input_checkpoint=ckpt.model_checkpoint_path,
output_node_names=target_nodes,
output_graph=(self.model_path + '.bytes'),
clear_devices=True, initializer_nodes='', input_saver='',
restore_op_name='save/restore_all',
filename_tensor_name='save/Const:0')
def main():
saver = tf.train.import_meta_graph(
'../datasets_test/cifar/model/model.ckpt-832.meta')
with tf.Session() as sess:
last_point = tf.train.latest_checkpoint(
'../datasets_test/cifar/model/')
saver.restore(sess, last_point)
tf.train.write_graph(sess.graph_def, 'output_model/pb_model',
'model.pb')
freeze_graph.freeze_graph('output_model/pb_model/model.pb', '', False,
model_path, 'add', 'save/restore_all',
'save/Const:0',
'output_model/pb_model/frozen_model.pb',
False, "")
def do_all(_):
if tf.gfile.Exists(FLAGS.checkpoint_dir):
tf.gfile.DeleteRecursively(FLAGS.checkpoint_dir)
tf.gfile.MakeDirs(FLAGS.checkpoint_dir)
if tf.gfile.Exists(FLAGS.model_dir):
tf.gfile.DeleteRecursively(FLAGS.model_dir)
tf.gfile.MakeDirs(FLAGS.model_dir)
content_targets = get_files(FLAGS.train_path)
print(FLAGS.style)
print(list_files(FLAGS.style))
for name in get_files(FLAGS.style):
print(name)
style_target = get_img(name)
name = name.split('/')[len(name.split('/'))-1]
name = name[:len(name)-4]
print(name)
for epoch, iteration, ckpt in optimize(
content_targets, style_target, content_weight,
style_weight, tv_weight, FLAGS.vgg_path, name, FLAGS.light_version):
if (FLAGS.test):
assert FLAGS.test_dir is not False
preds_img_name = "{}_{}.png".format(epoch, iteration)
preds_img_path = os.path.join(FLAGS.test_dir, preds_img_name)
evaluate(FLAGS.test, preds_img_path, ckpt, name, FLAGS.light_version)
start_time = time.time()
freeze_graph(
input_graph=os.path.join(FLAGS.model_dir,
FLAGS.model_name + '_' + name + '.pb.txt'),
input_saver='',
input_binary=False,
input_checkpoint=ckpt,
output_node_names='output',
restore_op_name='save/restore_all',
filename_tensor_name='save/Const:0',
output_graph=os.path.join(FLAGS.model_dir,
'%s_frozen.pb' % (FLAGS.model_name + '_' + name)),
clear_devices=False,
initializer_nodes='')
end_time = time.time()
ex_time = end_time - start_time
print('Save file graph frozen pb done!, time:', ex_time)
def main():
im_test = io.imread(
'/home/bo718.wang/irfan/pacling/train_2_out/images/u1025-inputs.png')
im_test = transform.resize(im_test, [512, 256, 3])
in_dir = "/home/bo718.wang/irfan/pacling/train_2/"
out_dir = "./out"
tf.reset_default_graph()
im = tf.placeholder(tf.float32, [None, None, 3])
# im1 = tf.placeholder(tf.float32, [None, None, 3])
# im1 = tf.placeholder(tf.float32, [None, None, 3])
im1 = tf.expand_dims(im, 0)
im1 *= 2
im1 -= 1
# im1 *= 256
# im1 -= 128
with tf.variable_scope('generator'):
g_out = create_generator(im1, 3)
# g_out *= 128
g_out /= 2
g_out += 0.5
# out = g_out + 128.
out = tf.identity(g_out, 'out')
# print(out.name)
# print('--------------------')
# print("type of out",type(g_out))
saver = tf.train.Saver(max_to_keep=1)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
pretrained = tf.train.latest_checkpoint(in_dir)
saver.restore(sess, pretrained)
out = sess.run(out, feed_dict={im: im_test})
io.imsave('out.jpg', out[0])
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
pretrained = tf.train.latest_checkpoint(in_dir)
saver.restore(sess, pretrained)
# TF SAVER
checkpoint_path = in_dir + 'model-881800'
# FREEZE GRAPH
tf.train.write_graph(sess.graph_def, in_dir, 'model.pb')
freeze_graph.freeze_graph(in_dir + '/model.pb', '', False,
checkpoint_path, 'out', 'save/restore_all',
'save/Const:0',
out_dir + '/119_cartoon_size_sketch.pb',
False, "")
print('==================')
def main():
# Define the name of your model
model_name = 'batch=100,lr=0.0001,optimizer=GDS,epochs=1000'
# define the path to the graph from training
input_graph = os.path.join(os.getcwd(), 'tensorflow_logs', model_name,
'graph', 'graph.pb')
# define the path in which to save the frozen graph
output_graph = os.path.join(os.getcwd(), 'tensorflow_logs', model_name,
'frozen_graph', 'frozen_graph.pb')
# the frozen_graph directory must exist in order to freeze the model
directory = os.path.join(os.getcwd(), 'tensorflow_logs', model_name,
'frozen_graph')
if not os.path.exists(directory):
os.makedirs(directory)
# define the checkpoint/weights you want to freeze inside the graph
input_checkpoint = os.path.join(os.getcwd(), 'tensorflow_logs', model_name,
'train-900')
# define the name of the prediction output node
# This name can be easily extracted using Tensorboard. In GRAPHS tab of Tensorboard, check the inputs of Loss scope.
# In this case they are "vel_true" and "ConvNet/fc_layer_2/BiasAdd".The CNN's predictions are provided from the
# "ConvNet/fc_layer_2/BiasAdd" element, whereas the true omega velocities from the "vel_true". Here we have to define
# the element which provides the CNN's predictions and thus we defined as output_node_names the "ConvNet/fc_layer_2/BiasAdd".
output_node_names = "ConvNet/fc_layer_2/BiasAdd"
# The following settings should remain the same
input_saver = ""
input_binary = True
restore_op_name = 'save/restore_all'
filename_tensor_name = 'save/Const:0'
clear_devices = True
initializer_nodes = ""
variable_names_blacklist = ""
# Freeze the graph
freeze_graph.freeze_graph(input_graph, input_saver, input_binary,
input_checkpoint, output_node_names,
restore_op_name, filename_tensor_name,
output_graph, clear_devices, initializer_nodes,
variable_names_blacklist)
print("The frozen graph is saved in {}.".format(output_graph))
def testFreezeSavedModel(self):
tmp_dir = self.get_temp_dir()
saved_model_dir = os.path.join(tmp_dir, "saved_model_dir")
feature_name = "feature"
self._writeDummySavedModel(saved_model_dir, feature_name)
output_graph_filename = os.path.join(tmp_dir, "output_graph.pb")
input_saved_model_dir = saved_model_dir
output_node_names = "output_node"
input_binary = False
input_saver_def_path = False
restore_op_name = None
filename_tensor_name = None
clear_devices = False
input_meta_graph = False
checkpoint_path = None
input_graph_filename = None
saved_model_tags = tag_constants.SERVING
freeze_graph.freeze_graph(input_graph_filename, input_saver_def_path,
input_binary, checkpoint_path,
output_node_names, restore_op_name,
filename_tensor_name, output_graph_filename,
clear_devices, "", "", "", input_meta_graph,
input_saved_model_dir, saved_model_tags)
# Now we make sure the variable is now a constant, and that the graph still
# produces the expected result.
with ops.Graph().as_default():
output_graph_def = graph_pb2.GraphDef()
with open(output_graph_filename, "rb") as f:
output_graph_def.ParseFromString(f.read())
_ = importer.import_graph_def(output_graph_def, name="")
self.assertEqual(8, len(output_graph_def.node))
for node in output_graph_def.node:
self.assertNotEqual("VariableV2", node.op)
self.assertNotEqual("Variable", node.op)
feature_value = 2.0
example = self._createTFExampleString(feature_name, feature_value)
with session.Session() as sess:
input_node = sess.graph.get_tensor_by_name("input_node:0")
output_node = sess.graph.get_tensor_by_name("output_node:0")
output = sess.run(output_node,
feed_dict={input_node: [example]})
self.assertNear(feature_value, output, 0.00001)
def freeze_saved_model_graph(output_node_names, input_saved_model_dir,
output_graph_filename):
input_binary = False
input_saver_def_path = False
restore_op_name = None
filename_tensor_name = None
clear_devices = False
input_meta_graph = False
checkpoint_path = None
input_graph_filename = None
saved_model_tags = tag_constants.SERVING
freeze_graph.freeze_graph(input_graph_filename, input_saver_def_path,
input_binary, checkpoint_path, output_node_names,
restore_op_name, filename_tensor_name,
output_graph_filename, clear_devices, "", "", "",
input_meta_graph, input_saved_model_dir,
saved_model_tags)
def export_model(input_node_names, output_node_name):
freeze_graph.freeze_graph('out/' + MODEL_NAME + '.pbtxt', None, False,
'out/' + MODEL_NAME + '.chkp', output_node_name, "save/restore_all",
"save/Const:0", 'out/frozen_' + MODEL_NAME + '.pb', True, "")
input_graph_def = tf.GraphDef()
with tf.gfile.Open('out/frozen_' + MODEL_NAME + '.pb', "rb") as f:
input_graph_def.ParseFromString(f.read())
output_graph_def = optimize_for_inference_lib.optimize_for_inference(
input_graph_def, input_node_names, [output_node_name],
tf.float32.as_datatype_enum)
with tf.gfile.FastGFile('out/opt_' + MODEL_NAME + '.pb', "wb") as f:
f.write(output_graph_def.SerializeToString())
print("graph saved!")
def _freeze_keras_saved_model(self, saved_model_dir):
"""Freezes the model and returns the frozen GraphDef.
Frozen here means that all variables are converted to placeholders.
Args:
saved_model_dir: Directory with the Keras SavedModel export.
Returns:
Frozen GraphDef for the model.
"""
temp_dir = tempfile.mkdtemp("tflite-transfer-convert")
graph_def_file_name = os.path.join(temp_dir, "frozen.pb")
output_names = [
utils.tensor_to_op_name(output.name)
for output in self._eval_signature.outputs.values()
]
freeze_graph.freeze_graph(
input_graph=None,
input_saver=False,
input_binary=True,
input_checkpoint=None,
output_node_names=",".join(output_names),
restore_op_name=None,
filename_tensor_name=None,
output_graph=graph_def_file_name,
clear_devices=True,
initializer_nodes="",
input_saved_model_dir=saved_model_dir,
saved_model_tags="eval",
)
const_graph_def = tfv1.GraphDef()
with open(graph_def_file_name, "rb") as graph_def_file:
const_graph_def.ParseFromString(graph_def_file.read())
# Convert constants produced from trainable variables to placeholders.
# Note: eval model might have other variables that should not be trainable,
# they are kept as constants. Only variables that are present in serve
# model are converted.
graph_def = utils.convert_constants_to_placeholders(
const_graph_def, self._variable_names)
shutil.rmtree(temp_dir)
return graph_def
def freeze_graph(sess, ckpt, output):
print("Loading checkpoint...")
saver = tf.train.Saver()
saver.restore(sess, ckpt)
print("Writing graph...")
if not os.path.isdir("_Cache"):
os.makedirs("_Cache")
_dir = os.path.join("_Cache", "Model")
saver.save(sess, _dir)
graph_io.write_graph(sess.graph, "_Cache", "Model.pb", False)
print("Freezing graph...")
freeze_graph.freeze_graph(
os.path.join("_Cache", "Model.pb"),
"", True, os.path.join("_Cache", "Model"),
output, "save/restore_all", "save/Const:0", "Frozen.pb", True, ""
)
print("Done")
def export_model(saver, input_node_names, output_node_name):
# creates the 'out' folder where our frozen graphs will be saved
if not path.exists('out'):
os.mkdir('out')
# an arbitrary name for our graph
GRAPH_NAME = 'heli_test_to _Unity'
# GRAPH SAVING - '.pbtxt'
tf.train.write_graph(K.get_session().graph_def, 'out',
GRAPH_NAME + '_graph.pbtxt')
# GRAPH SAVING - '.chkp'
# KEY: This method saves the graph at it's last checkpoint (hence '.chkp')
saver.save(K.get_session(), 'out/' + GRAPH_NAME + '.chkp')
# GRAPH SAVING - '.bytes'
# freeze_graph.freeze_graph(input_graph_path, input_saver_def_path,
# input_binary, checkpoint_path, output_node_names,
# restore_op_name, filename_tensor_name,
# output_frozen_graph_name, clear_devices, "")
freeze_graph.freeze_graph('out/' + GRAPH_NAME + '_graph.pbtxt', None,
False, 'out/' + GRAPH_NAME + '.chkp',
output_node_name, "save/restore_all",
"save/Const:0",
'out/frozen_' + GRAPH_NAME + '.bytes', True, "")
# freeze_graph.freeze_graph(input_graph='out/' + GRAPH_NAME + '_graph.pbtxt',
# input_binary=True,
# input_checkpoint='out/' + GRAPH_NAME + '.chkp',
# output_node_names=output_node_name,
# output_graph='out/frozen_' + GRAPH_NAME + '.bytes',
# clear_devices=True, initializer_nodes="", input_saver="",
# restore_op_name="save/restore_all", filename_tensor_name="save/Const:0")
# GRAPH OPTIMIZING
input_graph_def = tf.GraphDef()
with tf.gfile.Open('out/frozen_' + GRAPH_NAME + '.bytes', "rb") as f:
input_graph_def.ParseFromString(f.read())
output_graph_def = optimize_for_inference_lib.optimize_for_inference(
input_graph_def, input_node_names, [output_node_name],
tf.float32.as_datatype_enum)
with tf.gfile.FastGFile('out/opt_' + GRAPH_NAME + '.bytes', "wb") as f:
f.write(output_graph_def.SerializeToString())
print("graph saved!")
def freeze_model(model_path, freeze_dir, output_graph_name):
name = re.sub(r'-\d+$', '', os.path.basename(model_path))
model_dir = os.path.dirname(model_path)
model_type, _, _ = read_meta('%s/meta' % model_dir)
if not os.path.isdir(freeze_dir):
os.makedirs(freeze_dir)
copyfile('%s/meta' % model_dir, '%s/meta' % freeze_dir)
checkpoint_prefix = os.path.join(freeze_dir, "saved_checkpoint")
checkpoint_state_name = "checkpoint_state"
input_graph_name = "input_graph.pb"
G2PModel, hparams = import_model_type(model_type)
with open('%s/hparams' % model_dir, 'r') as infp:
loaded = json.load(infp)
hparams.parse_json(loaded)
with ops.Graph().as_default():
with tf.Session() as sess:
model = G2PModel(hparams, is_training=False, with_target=False, reuse=False)
saver = tf.train.Saver()
sess.run(tf.global_variables_initializer())
saver.restore(sess, model_path)
saver = saver_lib.Saver()
checkpoint_path = saver.save(
sess,
checkpoint_prefix,
global_step=0,
latest_filename=checkpoint_state_name)
graph_io.write_graph(sess.graph, freeze_dir, input_graph_name)
input_graph_path = os.path.join(freeze_dir, input_graph_name)
input_saver_def_path = ""
input_binary = False
output_node_names = 'g2p/predicted_1best,g2p/probs'
restore_op_name = "save/restore_all"
filename_tensor_name = "save/Const:0"
output_graph_path = os.path.join(freeze_dir, output_graph_name)
clear_devices = False
freeze_graph.freeze_graph(
input_graph_path, input_saver_def_path, input_binary, checkpoint_path,
output_node_names, restore_op_name, filename_tensor_name,
output_graph_path, clear_devices, "")
def convertmodeltopb2():
from tensorflow.python.tools import freeze_graph
output_node_names = ['StatefulPartitionedCall']
output_node_names = ','.join(output_node_names)
save_pb_model_path = './tfrec/fast_real_frozen3/faster_5x1_obj.pb'
input_saved_model_dir='./tfrec/fast_model_frozen/saved_model'
freeze_graph.freeze_graph(input_graph=None, input_saver=None,
input_binary=None,
input_checkpoint=None,
output_node_names=output_node_names,
restore_op_name=None,
filename_tensor_name=None,
output_graph=save_pb_model_path,
clear_devices=None,
initializer_nodes=None,
input_saved_model_dir=input_saved_model_dir)
def exportModelToTF(tfModelOutputDir):
if not os.path.exists(tfModelOutputDir):
os.makedirs(tfModelOutputDir)
# Save checkpoint
saver = tf.train.Saver()
save_path = saver.save(K.get_session(), tfModelOutputDir + "/model")
# Save metagraph
tf.train.write_graph(K.get_session().graph.as_graph_def(), "", tfModelOutputDir + "/metagraph.pb", False)
# Freeze graph
freeze_graph(input_graph=tfModelOutputDir + "/metagraph.pb", input_saver="", input_binary=True,
input_checkpoint=tfModelOutputDir + "/model", output_node_names='softmax/Softmax',
restore_op_name="save/restore_all", filename_tensor_name="save/Const:0",
output_graph=tfModelOutputDir + "/graph.pb", clear_devices=True, initializer_nodes="")
def convert_keras_to_freeze_pb(model, frozen_model_path):
out_names = ",".join([layer.name.split(":")[0] for layer in model.outputs])
inp_names = ",".join([layer.name.split(":")[0] for layer in model.inputs])
print("OUTPUT: {}".format(out_names))
print("INPUT: {}".format(inp_names))
model.summary()
K.set_learning_phase(0)
sess = K.get_session()
saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V2)
checkpoint_path = saver.save(sess, './saved_ckpt', global_step=0, latest_filename='checkpoint_state')
graph_io.write_graph(sess.graph, '.', './tmp.pb')
freeze_graph.freeze_graph('./tmp.pb', '',
False, checkpoint_path, out_names,
"save/restore_all",
"save/Const:0",
frozen_model_path,
False, "")
def testFreezeSavedModel(self):
tmp_dir = self.get_temp_dir()
saved_model_dir = os.path.join(tmp_dir, "saved_model_dir")
feature_name = "feature"
self._writeDummySavedModel(saved_model_dir, feature_name)
output_graph_filename = os.path.join(tmp_dir, "output_graph.pb")
input_saved_model_dir = saved_model_dir
output_node_names = "output_node"
input_binary = False
input_saver_def_path = False
restore_op_name = None
filename_tensor_name = None
clear_devices = False
input_meta_graph = False
checkpoint_path = None
input_graph_filename = None
saved_model_tags = tag_constants.SERVING
freeze_graph.freeze_graph(input_graph_filename, input_saver_def_path,
input_binary, checkpoint_path, output_node_names,
restore_op_name, filename_tensor_name,
output_graph_filename, clear_devices, "", "", "",
input_meta_graph, input_saved_model_dir,
saved_model_tags)
# Now we make sure the variable is now a constant, and that the graph still
# produces the expected result.
with ops.Graph().as_default():
output_graph_def = graph_pb2.GraphDef()
with open(output_graph_filename, "rb") as f:
output_graph_def.ParseFromString(f.read())
_ = importer.import_graph_def(output_graph_def, name="")
self.assertEqual(8, len(output_graph_def.node))
for node in output_graph_def.node:
self.assertNotEqual("VariableV2", node.op)
self.assertNotEqual("Variable", node.op)
feature_value = 2.0
example = self._createTFExampleString(feature_name, feature_value)
with session.Session() as sess:
input_node = sess.graph.get_tensor_by_name("input_node:0")
output_node = sess.graph.get_tensor_by_name("output_node:0")
output = sess.run(output_node, feed_dict={input_node: [example]})
self.assertNear(feature_value, output, 0.00001)
def convert_to_frozen_graph(self):
input_pb_path = self.frozenpb_config_worker.input_dir_path + self.frozenpb_config_worker.input_pb_name
input_ckpt_path = self.frozenpb_config_worker.input_dir_path + self.frozenpb_config_worker.input_ckpt_name
output_frozen_pb_path = self.frozenpb_config_worker.output_dir_path + self.frozenpb_config_worker.output_pb_name
freeze_graph.freeze_graph(\
input_graph=input_pb_path,
input_saver= "", # this argument is used with SavedModel
input_binary=self.frozenpb_config_worker.binary_opt,
input_checkpoint=input_ckpt_path,
output_node_names=self.frozenpb_config_worker.output_node_names,
restore_op_name="save/restore_all", # unused in freeze_graph()
filename_tensor_name="save/Const:0", # unused in freeze_graph()
output_graph=output_frozen_pb_path,
clear_devices=False, # not clear how to use
initializer_nodes="")
def _load_saved_model(self):
"""Load the tensorflow saved model."""
try:
from tensorflow.python.tools import freeze_graph
from tensorflow.python.framework import ops
from tensorflow.python.framework import graph_util
except ImportError:
raise ImportError(
"InputConfiguration: Unable to import tensorflow which is "
"required to restore from saved model.")
saved_model_dir = self._model_dir
output_graph_filename = self._tmp_dir.relpath("tf_frozen_model.pb")
input_saved_model_dir = saved_model_dir
output_node_names = self._get_output_names()
input_binary = False
input_saver_def_path = False
restore_op_name = None
filename_tensor_name = None
clear_devices = True
input_meta_graph = False
checkpoint_path = None
input_graph_filename = None
saved_model_tags = ",".join(self._get_tag_set())
freeze_graph.freeze_graph(input_graph_filename, input_saver_def_path,
input_binary, checkpoint_path, output_node_names,
restore_op_name, filename_tensor_name,
output_graph_filename, clear_devices, "", "", "",
input_meta_graph, input_saved_model_dir,
saved_model_tags)
with ops.Graph().as_default():
output_graph_def = graph_pb2.GraphDef()
with open(output_graph_filename, "rb") as f:
output_graph_def.ParseFromString(f.read())
output_graph_def = graph_util.remove_training_nodes(output_graph_def)
return output_graph_def
def saveModel(self, sess, outputDirectory = ""):
from tensorflow.python.framework import graph_io
from tensorflow.python.tools import freeze_graph
input_graph_path = outputDirectory + "tfModel.pb"
graph_io.write_graph(sess.graph, "./", input_graph_path)
#create frozen version of graph for distribution
input_saver_def_path = ""
input_binary = False
checkpoint_path = outputDirectory + "models/model.ckpt"
output_node_names = "y_ph"
restore_op_name = "save/restore_all"
filename_tensor_name = "save/Const:0"
output_graph_path = outputDirectory + "tfModel_frozen.pb"
clear_devices = False
freeze_graph.freeze_graph(input_graph_path, input_saver_def_path,
input_binary, checkpoint_path, output_node_names,
restore_op_name, filename_tensor_name,
output_graph_path, clear_devices, "")
print("Frozen model (model and weights) saved in file: %s" % output_graph_path)
def dump_frozen_graph(sess, graph_file, output_node_names=None):
assert graph_file.endswith('.pb')
assert output_node_names is None or isinstance(output_node_names, list)
output_node_names = output_node_names or estimate_inputs_outputs(sess.graph)[1]
dir_ = os.path.dirname(graph_file)
base = os.path.basename(graph_file)
ckpt = graph_file.replace('.pb', '.ckpt')
frozen = graph_file.replace('.pb', '.pb.frozen')
os.system('mkdir -p {}'.format(dir_))
print('>> Saving `{}`... '.format(graph_file), end='')
tf.train.write_graph(sess.graph, dir_, base, as_text=False)
tf.train.write_graph(sess.graph, dir_, base + "txt", as_text=True)
print('Done')
print('>> Saving `{}`... '.format(ckpt), end='')
tf.train.Saver().save(sess, ckpt, write_meta_graph=False)
print('Done')
print('>> Freezing graph to `{}`... '.format(frozen))
print('Outputs:\n {}'.format(', '.join(output_node_names)))
from tensorflow.python.tools.freeze_graph import freeze_graph
freeze_graph(input_graph=graph_file,
input_saver='',
input_binary=True,
input_checkpoint=ckpt,
output_node_names=','.join(output_node_names),
restore_op_name='save/restore_all',
filename_tensor_name='save/Const:0',
output_graph=frozen,
clear_devices=True,
initializer_nodes='',
saved_model_tags='serve')
return frozen
def main(_):
# Always keep the cpu as default
with tf.Graph().as_default(), tf.device('/cpu:0'):
if FLAGS.validation_interval == 0:
FLAGS.validation_db = None
# Set Tensorboard log directory
if FLAGS.summaries_dir:
# The following gives a nice but unrobust timestamp
FLAGS.summaries_dir = os.path.join(FLAGS.summaries_dir, datetime.datetime.now().strftime("%Y%m%d_%H%M%S"))
if not FLAGS.train_db and not FLAGS.validation_db and not FLAGS.inference_db and not FLAGS.visualizeModelPath:
logging.error("At least one of the following file sources should be specified: "
"train_db, validation_db or inference_db")
exit(-1)
if FLAGS.seed:
tf.set_random_seed(FLAGS.seed)
batch_size_train = FLAGS.batch_size
batch_size_val = FLAGS.batch_size
logging.info("Train batch size is %s and validation batch size is %s", batch_size_train, batch_size_val)
# This variable keeps track of next epoch, when to perform validation.
next_validation = FLAGS.validation_interval
logging.info("Training epochs to be completed for each validation : %s", next_validation)
# This variable keeps track of next epoch, when to save model weights.
next_snapshot_save = FLAGS.snapshotInterval
logging.info("Training epochs to be completed before taking a snapshot : %s", next_snapshot_save)
last_snapshot_save_epoch = 0
snapshot_prefix = FLAGS.snapshotPrefix if FLAGS.snapshotPrefix else FLAGS.network.split('.')[0]
logging.info("Model weights will be saved as %s_<EPOCH>_Model.ckpt", snapshot_prefix)
if not os.path.exists(FLAGS.save):
os.makedirs(FLAGS.save)
logging.info("Created a directory %s to save all the snapshots", FLAGS.save)
# Load mean variable
if FLAGS.subtractMean == 'none':
mean_loader = None
else:
if not FLAGS.mean:
logging.error("subtractMean parameter not set to 'none' yet mean image path is unset")
exit(-1)
logging.info("Loading mean tensor from %s file", FLAGS.mean)
mean_loader = tf_data.MeanLoader(FLAGS.mean, FLAGS.subtractMean, FLAGS.bitdepth)
classes = 0
nclasses = 0
if FLAGS.labels_list:
logging.info("Loading label definitions from %s file", FLAGS.labels_list)
classes = loadLabels(FLAGS.labels_list)
nclasses = len(classes)
if not classes:
logging.error("Reading labels file %s failed.", FLAGS.labels_list)
exit(-1)
logging.info("Found %s classes", nclasses)
# Create a data-augmentation dict
aug_dict = {
'aug_flip': FLAGS.augFlip,
'aug_noise': FLAGS.augNoise,
'aug_contrast': FLAGS.augContrast,
'aug_whitening': FLAGS.augWhitening,
'aug_HSV': {
'h': FLAGS.augHSVh,
's': FLAGS.augHSVs,
'v': FLAGS.augHSVv,
},
}
# Import the network file
path_network = os.path.join(os.path.dirname(os.path.realpath(__file__)), FLAGS.networkDirectory, FLAGS.network)
exec(open(path_network).read(), globals())
try:
UserModel
except NameError:
logging.error("The user model class 'UserModel' is not defined.")
exit(-1)
if not inspect.isclass(UserModel): # noqa
logging.error("The user model class 'UserModel' is not a class.")
exit(-1)
# @TODO(tzaman) - add mode checks to UserModel
if FLAGS.train_db:
with tf.name_scope(digits.STAGE_TRAIN) as stage_scope:
train_model = Model(digits.STAGE_TRAIN, FLAGS.croplen, nclasses, FLAGS.optimization, FLAGS.momentum)
train_model.create_dataloader(FLAGS.train_db)
train_model.dataloader.setup(FLAGS.train_labels,
FLAGS.shuffle,
FLAGS.bitdepth,
batch_size_train,
FLAGS.epoch,
FLAGS.seed)
train_model.dataloader.set_augmentation(mean_loader, aug_dict)
train_model.create_model(UserModel, stage_scope) # noqa
if FLAGS.validation_db:
with tf.name_scope(digits.STAGE_VAL) as stage_scope:
val_model = Model(digits.STAGE_VAL, FLAGS.croplen, nclasses, reuse_variable=True)
val_model.create_dataloader(FLAGS.validation_db)
val_model.dataloader.setup(FLAGS.validation_labels,
False,
FLAGS.bitdepth,
batch_size_val,
1e9,
FLAGS.seed) # @TODO(tzaman): set numepochs to 1
val_model.dataloader.set_augmentation(mean_loader)
val_model.create_model(UserModel, stage_scope) # noqa
if FLAGS.inference_db:
with tf.name_scope(digits.STAGE_INF) as stage_scope:
inf_model = Model(digits.STAGE_INF, FLAGS.croplen, nclasses)
inf_model.create_dataloader(FLAGS.inference_db)
inf_model.dataloader.setup(None, False, FLAGS.bitdepth, FLAGS.batch_size, 1, FLAGS.seed)
inf_model.dataloader.set_augmentation(mean_loader)
inf_model.create_model(UserModel, stage_scope) # noqa
# Start running operations on the Graph. allow_soft_placement must be set to
# True to build towers on GPU, as some of the ops do not have GPU
# implementations.
sess = tf.Session(config=tf.ConfigProto(
allow_soft_placement=True, # will automatically do non-gpu supported ops on cpu
inter_op_parallelism_threads=TF_INTER_OP_THREADS,
intra_op_parallelism_threads=TF_INTRA_OP_THREADS,
log_device_placement=FLAGS.log_device_placement))
if FLAGS.visualizeModelPath:
visualize_graph(sess.graph_def, FLAGS.visualizeModelPath)
exit(0)
# Saver creation.
if FLAGS.save_vars == 'all':
vars_to_save = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
elif FLAGS.save_vars == 'trainable':
vars_to_save = tf.all_variables()
else:
logging.error('Unknown save_var flag (%s)' % FLAGS.save_vars)
exit(-1)
saver = tf.train.Saver(vars_to_save, max_to_keep=0, sharded=FLAGS.serving_export)
# Initialize variables
init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
sess.run(init_op)
# If weights option is set, preload weights from existing models appropriately
if FLAGS.weights:
load_snapshot(sess, FLAGS.weights, tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES))
# Tensorboard: Merge all the summaries and write them out
writer = tf.summary.FileWriter(os.path.join(FLAGS.summaries_dir, 'tb'), sess.graph)
# If we are inferencing, only do that.
if FLAGS.inference_db:
inf_model.start_queue_runners(sess)
Inference(sess, inf_model)
queue_size_op = []
for n in tf.get_default_graph().as_graph_def().node:
if '_Size' in n.name:
queue_size_op.append(n.name+':0')
start = time.time() # @TODO(tzaman) - removeme
# Initial Forward Validation Pass
if FLAGS.validation_db:
val_model.start_queue_runners(sess)
Validation(sess, val_model, 0)
if FLAGS.train_db:
# During training, a log output should occur at least X times per epoch or every X images, whichever lower
train_steps_per_epoch = train_model.dataloader.get_total() / batch_size_train
if math.ceil(train_steps_per_epoch/MIN_LOGS_PER_TRAIN_EPOCH) < math.ceil(5000/batch_size_train):
logging_interval_step = int(math.ceil(train_steps_per_epoch/MIN_LOGS_PER_TRAIN_EPOCH))
else:
logging_interval_step = int(math.ceil(5000/batch_size_train))
logging.info("During training. details will be logged after every %s steps (batches)",
logging_interval_step)
# epoch value will be calculated for every batch size. To maintain unique epoch value between batches,
# it needs to be rounded to the required number of significant digits.
epoch_round = 0 # holds the required number of significant digits for round function.
tmp_batchsize = batch_size_train*logging_interval_step
while tmp_batchsize <= train_model.dataloader.get_total():
tmp_batchsize = tmp_batchsize * 10
epoch_round += 1
logging.info("While logging, epoch value will be rounded to %s significant digits", epoch_round)
# Create the learning rate policy
total_training_steps = train_model.dataloader.num_epochs * train_model.dataloader.get_total() / \
train_model.dataloader.batch_size
lrpolicy = lr_policy.LRPolicy(FLAGS.lr_policy,
FLAGS.lr_base_rate,
FLAGS.lr_gamma,
FLAGS.lr_power,
total_training_steps,
FLAGS.lr_stepvalues)
train_model.start_queue_runners(sess)
# Training
logging.info('Started training the model')
current_epoch = 0
try:
step = 0
step_last_log = 0
print_vals_sum = 0
while not train_model.queue_coord.should_stop():
log_runtime = FLAGS.log_runtime_stats_per_step and (step % FLAGS.log_runtime_stats_per_step == 0)
run_options = None
run_metadata = None
if log_runtime:
# For a HARDWARE_TRACE you need NVIDIA CUPTI, a 'CUDA-EXTRA'
# SOFTWARE_TRACE HARDWARE_TRACE FULL_TRACE
run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
feed_dict = {train_model.learning_rate: lrpolicy.get_learning_rate(step)}
if False:
for op in train_model.train:
_, summary_str, step = sess.run([op, train_model.summary, train_model.global_step],
feed_dict=feed_dict,
options=run_options,
run_metadata=run_metadata)
else:
_, summary_str, step = sess.run([train_model.train,
train_model.summary,
train_model.global_step],
feed_dict=feed_dict,
options=run_options,
run_metadata=run_metadata)
# HACK
step = step / len(train_model.train)
# logging.info(sess.run(queue_size_op)) # DEVELOPMENT: for checking the queue size
if log_runtime:
writer.add_run_metadata(run_metadata, str(step))
save_timeline_trace(run_metadata, FLAGS.save, int(step))
writer.add_summary(summary_str, step)
# Parse the summary
tags, print_vals = summary_to_lists(summary_str)
print_vals_sum = print_vals + print_vals_sum
# @TODO(tzaman): account for variable batch_size value on very last epoch
current_epoch = round((step * batch_size_train) / train_model.dataloader.get_total(), epoch_round)
# Start with a forward pass
if ((step % logging_interval_step) == 0):
steps_since_log = step - step_last_log
print_list = print_summarylist(tags, print_vals_sum/steps_since_log)
logging.info("Training (epoch " + str(current_epoch) + "): " + print_list)
print_vals_sum = 0
step_last_log = step
# Potential Validation Pass
if FLAGS.validation_db and current_epoch >= next_validation:
Validation(sess, val_model, current_epoch)
# Find next nearest epoch value that exactly divisible by FLAGS.validation_interval:
next_validation = (round(float(current_epoch)/FLAGS.validation_interval) + 1) * \
FLAGS.validation_interval
# Saving Snapshot
if FLAGS.snapshotInterval > 0 and current_epoch >= next_snapshot_save:
checkpoint_path, graphdef_path = save_snapshot(sess,
saver,
FLAGS.save,
snapshot_prefix,
current_epoch,
FLAGS.serving_export
)
# To find next nearest epoch value that exactly divisible by FLAGS.snapshotInterval
next_snapshot_save = (round(float(current_epoch)/FLAGS.snapshotInterval) + 1) * \
FLAGS.snapshotInterval
last_snapshot_save_epoch = current_epoch
writer.flush()
except tf.errors.OutOfRangeError:
logging.info('Done training for epochs: tf.errors.OutOfRangeError')
except ValueError as err:
logging.error(err.args[0])
exit(-1) # DIGITS wants a dirty error.
except (KeyboardInterrupt):
logging.info('Interrupt signal received.')
# If required, perform final snapshot save
if FLAGS.snapshotInterval > 0 and FLAGS.epoch > last_snapshot_save_epoch:
checkpoint_path, graphdef_path =\
save_snapshot(sess, saver, FLAGS.save, snapshot_prefix, FLAGS.epoch, FLAGS.serving_export)
print('Training wall-time:', time.time()-start) # @TODO(tzaman) - removeme
# If required, perform final Validation pass
if FLAGS.validation_db and current_epoch >= next_validation:
Validation(sess, val_model, current_epoch)
if FLAGS.train_db:
if FLAGS.labels_list:
output_tensor = train_model.towers[0].inference
out_name, _ = output_tensor.name.split(':')
if FLAGS.train_db:
del train_model
if FLAGS.validation_db:
del val_model
if FLAGS.inference_db:
del inf_model
# We need to call sess.close() because we've used a with block
sess.close()
writer.close()
tf.reset_default_graph()
del sess
if FLAGS.train_db:
if FLAGS.labels_list:
path_frozen = os.path.join(FLAGS.save, 'frozen_model.pb')
print('Saving frozen model at path {}'.format(path_frozen))
freeze_graph.freeze_graph(
input_graph=graphdef_path,
input_saver='',
input_binary=True,
input_checkpoint=checkpoint_path,
output_node_names=out_name,
restore_op_name="save/restore_all",
filename_tensor_name="save/Const:0",
output_graph=path_frozen,
clear_devices=True,
initializer_nodes="",
)
logging.info('END')
exit(0)
try:
model= Sequential.from_config(config)
except:
model= Model.from_config(config)
#model= model_from_config(config)
model.set_weights(weights)
model.summary()
print("Input name:")
print(model.input.name)
print("Output name:")
print(model.output.name)
output_name=model.output.name.split(':')[0]
# not sure what this is for
export_version = 1 # version number (integer)
graph_file=export_path+".pb"
ckpt_file=export_path+".ckpt"
# create a saver
saver = tf.train.Saver(sharded=True)
tf.train.write_graph(sess.graph_def, '', graph_file)
save_path = saver.save(sess, ckpt_file)
input_graph_path = 'export/model.pb'
checkpoint_path = 'export/model.ckpt'
freeze_graph.freeze_graph(input_graph_path, "", False, checkpoint_path, "output_1/Softmax", "save/restore_all", "save/Const", output_frozen_graph_name, True, "")
def save(self, path=None, name=None, overwrite=True):
path = "Models" if path is None else path
name = "Cache" if name is None else name
folder = os.path.join(path, name)
if not os.path.exists(folder):
os.makedirs(folder)
_dir = os.path.join(folder, "Model")
if os.path.isfile(_dir):
if not overwrite:
_count = 1
_new_dir = _dir + "({})".format(_count)
while os.path.isfile(_new_dir):
_count += 1
_new_dir = _dir + "({})".format(_count)
_dir = _new_dir
else:
os.remove(_dir)
print()
print("=" * 60)
print("Saving Model to {}...".format(folder))
print("-" * 60)
with open(_dir + ".nn", "wb") as file:
# We don't need w_stds & b_inits when we load a model
_dic = {
"structures": {
"_lr": self._lr,
"_layer_names": self.layer_names,
"_layer_params": self._layer_params,
"_next_dimension": self._current_dimension
},
"params": {
"_logs": self._logs,
"_metric_names": self._metric_names,
"_optimizer": self._optimizer.name,
"layer_special_params": self.layer_special_params
}
}
pickle.dump(_dic, file)
saver = tf.train.Saver()
saver.save(self._sess, _dir)
graph_io.write_graph(self._sess.graph, os.path.join(path, name), "Model.pb", False)
with tf.name_scope("OutputFlow"):
self.get_rs(self._tfx)
_output = ""
for op in self._sess.graph.get_operations()[::-1]:
if "OutputFlow" in op.name:
_output = op.name
break
with open(os.path.join(path, name, "IO.txt"), "w") as file:
file.write("\n".join([
"Input : Entry/Placeholder:0",
"Output : {}:0".format(_output)
]))
graph_io.write_graph(self._sess.graph, os.path.join(path, name), "Cache.pb", False)
freeze_graph.freeze_graph(
os.path.join(path, name, "Cache.pb"),
"", True, os.path.join(path, name, "Model"),
_output, "save/restore_all", "save/Const:0",
os.path.join(path, name, "Frozen.pb"), True, ""
)
os.remove(os.path.join(path, name, "Cache.pb"))
print("Done")
print("=" * 60)
def single_worker_inference(infer_model,
ckpt,
inference_input_file,
inference_output_file,
hparams):
"""Inference with a single worker."""
output_infer = inference_output_file
# Read data
infer_data = load_data(inference_input_file, hparams)
print ("Batch size type:", type(hparams.infer_batch_size))
with tf.Session(
graph=infer_model.graph, config=utils.get_config_proto()) as sess:
# revo debug
# sess = tf_debug.TensorBoardDebugWrapperSession(sess, 'xy:6064')
# initi table
# sess.run(infer_model.insert_op[0])
# sess.run(infer_model.insert_op[1])
# sess.run(infer_model.insert_op[2])
#
loaded_infer_model = model_helper.load_model(
infer_model.model, ckpt, sess, "infer", infer_model.insert_op)
sess.run(
infer_model.iterator.initializer,
feed_dict={
infer_model.src_placeholder: infer_data,
infer_model.batch_size_placeholder: hparams.infer_batch_size
})
# Debug By Revo
# value = sess.run(infer_model.iterator.source)
# print ("Value:", value)
# print ("Value,len:", len(value))
# print ("Value,Type:", type(value))
# print ("Value,Shape:", value.shape)
# tmp_i = sess.run(infer_model.iterator)
# print ("iterator:", tmp_i)
# print ("iterator shape:", tmp_i.shape())
# sys.exit()
# print ("TEST")
# # Initialize keys and values.
# keys = tf.constant([1, 2, 3], dtype=tf.int64)
# vals = tf.constant([1, 2, 3], dtype=tf.int64)
# # Initialize hash table.
# table = tf.contrib.lookup.MutableDenseHashTable(key_dtype=tf.int64, value_dtype=tf.int64, default_value=-1,
# empty_key=0)
# # Insert values to hash table and run the op.
# insert_op = table.insert(keys, vals)
# sess.run(insert_op)
# # Print hash table lookups.
# print(sess.run(table.lookup(keys)))
# print("HERE2")
# Saving Decoder model
# Decode
utils.print_out("# Start decoding ff3")
# print ("indices:", hparams.inference_indices)
if hparams.inference_indices:
_decode_inference_indices(
loaded_infer_model,
sess,
output_infer=output_infer,
output_infer_summary_prefix=output_infer,
inference_indices=hparams.inference_indices,
tgt_eos=hparams.eos,
subword_option=hparams.subword_option)
else:
nmt_utils.decode_and_evaluate(
"infer",
loaded_infer_model,
sess,
output_infer,
ref_file=None,
metrics=hparams.metrics,
subword_option=hparams.subword_option,
beam_width=hparams.beam_width,
tgt_eos=hparams.eos,
num_translations_per_input=hparams.num_translations_per_input)
# saving model
OUTPUT_FOLDER = '7.19'
utils.print_out("Ouput Folder : " + OUTPUT_FOLDER)
utils.print_out("# Saving Decoder model (Normal,ckpt) By Revo")
loaded_infer_model.saver.save(sess, OUTPUT_FOLDER+"/current.ckpt")
# save pb file
graph_io.write_graph(sess.graph_def, OUTPUT_FOLDER, "current.graphdef")
tf.train.export_meta_graph(filename=OUTPUT_FOLDER + '/current.meta')
writer = tf.summary.FileWriter(OUTPUT_FOLDER, sess.graph)
writer.close()
# Frozen graph saving
OUTPUT_FROZEN_FILE = 'nmt.pb'
# OUTPUT_NAMES = ['index_to_string_Lookup', 'table_init', 'batch_iter_init']
# maybe it is not utf8 as output
OUTPUT_NODES = ['reverse_table_Lookup']
utils.print_out("# Saving Decoder model (Frozen) By Revo")
# extract method try
# new_graph_def = tf.graph_util.extract_sub_graph(sess.graph_def, ["hash_table_2_Lookup"])
#
# remove train node
utils.print_out("# Removed Training nodes and outputing graph_def")
inference_graph = tf.graph_util.remove_training_nodes(sess.graph.as_graph_def())
graph_io.write_graph(inference_graph, OUTPUT_FOLDER, "infer_model.graphdef")
# This is oLd version freeze graph
freeze_graph.freeze_graph("7.19/current.graphdef", "", False, "7.19/current.ckpt", "reverse_table_Lookup", "", "", OUTPUT_FOLDER + "/" + OUTPUT_FROZEN_FILE, True, "")
#
# frozen_graph = tf.graph_util.convert_variables_to_constants(sess, inference_graph, OUTPUT_NODES)
# Normal
frozen_graph = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, OUTPUT_NODES)
# graph_io.write_graph(frozen_graph.graph_def, OUTPUT_FOLDER, "frozen.graphdef")
# frozen_graph = tf.graph_util.convert_variables_to_constants(sess, new_graph_def, OUTPUT_NODES)
#
# tf.train.write_graph(frozen_graph, OUTPUT_FOLDER, OUTPUT_FROZEN_FILE, as_text=False)
# TOCO with python
utils.print_out("# Start converting into TOCO file.")
converter = tf.contrib.lite.TocoConverter.from_frozen_graph(OUTPUT_FOLDER + "/" + OUTPUT_FROZEN_FILE, ['src_place'], OUTPUT_NODES)
# try session way
# input = sess.graph.get_tensor_by_name("src_place:0")
# output = sess.graph.get_tensor_by_name("reverse_table_Lookup:0")
# converter = tf.contrib.lite.TocoConverter.from_session(sess, [input], [output])
#
tflite_model = converter.convert()
open(OUTPUT_FOLDER + "/converted_model.tflite", "wb").write(tflite_model)
def main(args):
# Read model parameters
checkpoint_path = tf.train.latest_checkpoint(args.checkpoint_dir)
if checkpoint_path is None:
log.error('Could not find a checkpoint in {}'.format(args.checkpoint_dir))
return
metapath = ".".join([checkpoint_path, "meta"])
log.info("Loading {}".format(metapath))
tf.train.import_meta_graph(metapath)
with tf.Session() as sess:
model_params = utils.get_model_params(sess)
if not hasattr(models, model_params['model_name']):
log.error("Model {} does not exist".format(model_params['model_name']))
return
mdl = getattr(models, model_params['model_name'])
# Instantiate new evaluation graph
tf.reset_default_graph()
sz = model_params['net_input_size']
log.info("Model {}".format(model_params['model_name']))
input_tensor = tf.placeholder(tf.float32, [1, sz, sz, 3], name='lowres_input')
with tf.variable_scope('inference'):
prediction = mdl.inference(input_tensor, input_tensor, model_params, is_training=False)
if model_params["model_name" ] == "HDRNetGaussianPyrNN":
output_tensor = tf.get_collection('packed_coefficients')[0]
output_tensor = tf.transpose(tf.squeeze(output_tensor), [3, 2, 0, 1, 4], name="output_coefficients")
log.info("Output shape".format(output_tensor.get_shape()))
else:
output_tensor = tf.get_collection('packed_coefficients')[0]
output_tensor = tf.transpose(tf.squeeze(output_tensor), [3, 2, 0, 1, 4], name="output_coefficients")
log.info("Output shape {}".format(output_tensor.get_shape()))
saver = tf.train.Saver()
gdef = tf.get_default_graph().as_graph_def()
log.info("Restoring weights from {}".format(checkpoint_path))
test_graph_name = "test_graph.pbtxt"
with tf.Session() as sess:
saver.restore(sess, checkpoint_path)
tf.train.write_graph(sess.graph, args.checkpoint_dir, test_graph_name)
input_graph_path = os.path.join(args.checkpoint_dir, test_graph_name)
output_graph_path = os.path.join(args.checkpoint_dir, "frozen_graph.pb")
input_saver_def_path = ""
input_binary = False
output_binary = True
input_node_names = input_tensor.name.split(":")[0]
output_node_names = output_tensor.name.split(":")[0]
restore_op_name = "save/restore_all"
filename_tensor_name = "save/Const:0"
clear_devices = False
log.info("Freezing to {}".format(output_graph_path))
freeze_graph.freeze_graph(input_graph_path, input_saver_def_path,
input_binary, checkpoint_path, output_node_names,
restore_op_name, filename_tensor_name,
output_graph_path, clear_devices, "")
log.info('input tensor: {} {}'.format(input_tensor.name, input_tensor.shape))
log.info('output tensor: {} {}'.format(output_tensor.name, output_tensor.shape))
# Dump guide parameters
if model_params['model_name'] == 'HDRNetCurves':
g = tf.get_default_graph()
ccm = g.get_tensor_by_name('inference/guide/ccm:0')
ccm_bias = g.get_tensor_by_name('inference/guide/ccm_bias:0')
shifts = g.get_tensor_by_name('inference/guide/shifts:0')
slopes = g.get_tensor_by_name('inference/guide/slopes:0')
mixing_weights = g.get_tensor_by_name('inference/guide/channel_mixing/weights:0')
mixing_bias = g.get_tensor_by_name('inference/guide/channel_mixing/biases:0')
ccm_, ccm_bias_, shifts_, slopes_, mixing_weights_, mixing_bias_ = sess.run(
[ccm, ccm_bias, shifts, slopes, mixing_weights, mixing_bias])
shifts_ = np.squeeze(shifts_).astype(np.float32)
slopes_ = np.squeeze(slopes_).astype(np.float32)
mix_matrix_dump = np.append(np.squeeze(mixing_weights_), mixing_bias_[0]).astype(np.float32)
ccm34_ = np.vstack((ccm_, ccm_bias_[np.newaxis, :]))
save(ccm34_.T, os.path.join(args.checkpoint_dir, 'guide_ccm_f32_3x4.bin'))
save(shifts_.T, os.path.join(args.checkpoint_dir, 'guide_shifts_f32_16x3.bin'))
save(slopes_.T, os.path.join(args.checkpoint_dir, 'guide_slopes_f32_16x3.bin'))
save(mix_matrix_dump, os.path.join(args.checkpoint_dir, 'guide_mix_matrix_f32_1x4.bin'))
elif model_params['model_name'] == "HDRNetGaussianPyrNN":
g = tf.get_default_graph()
for lvl in range(3):
conv1_w = g.get_tensor_by_name('inference/guide/level_{}/conv1/weights:0'.format(lvl))
conv1_b = g.get_tensor_by_name('inference/guide/level_{}/conv1/BatchNorm/beta:0'.format(lvl))
conv1_mu = g.get_tensor_by_name('inference/guide/level_{}/conv1/BatchNorm/moving_mean:0'.format(lvl))
conv1_sigma = g.get_tensor_by_name('inference/guide/level_{}/conv1/BatchNorm/moving_variance:0'.format(lvl))
conv1_eps = g.get_tensor_by_name('inference/guide/level_{}/conv1/BatchNorm/batchnorm/add/y:0'.format(lvl))
conv2_w = g.get_tensor_by_name('inference/guide/level_{}/conv2/weights:0'.format(lvl))
conv2_b = g.get_tensor_by_name('inference/guide/level_{}/conv2/biases:0'.format(lvl))
conv1w_, conv1b_, conv1mu_, conv1sigma_, conv1eps_, conv2w_, conv2b_ = sess.run(
[conv1_w, conv1_b, conv1_mu, conv1_sigma, conv1_eps, conv2_w, conv2_b])
conv1b_ -= conv1mu_/np.sqrt((conv1sigma_+conv1eps_))
conv1w_ = conv1w_/np.sqrt((conv1sigma_+conv1eps_))
conv1w_ = np.squeeze(conv1w_.astype(np.float32))
conv1b_ = np.squeeze(conv1b_.astype(np.float32))
conv1b_ = conv1b_[np.newaxis, :]
conv2w_ = np.squeeze(conv2w_.astype(np.float32))
conv2b_ = np.squeeze(conv2b_.astype(np.float32))
conv2 = np.append(conv2w_, conv2b_)
conv1 = np.vstack([conv1w_, conv1b_])
save(conv1.T, os.path.join(args.checkpoint_dir, 'guide_level{}_conv1.bin'.format(lvl)))
save(conv2, os.path.join(args.checkpoint_dir, 'guide_level{}_conv2.bin'.format(lvl)))
elif model_params['model_name'] in "HDRNetPointwiseNNGuide":
g = tf.get_default_graph()
conv1_w = g.get_tensor_by_name('inference/guide/conv1/weights:0')
conv1_b = g.get_tensor_by_name('inference/guide/conv1/BatchNorm/beta:0')
conv1_mu = g.get_tensor_by_name('inference/guide/conv1/BatchNorm/moving_mean:0')
conv1_sigma = g.get_tensor_by_name('inference/guide/conv1/BatchNorm/moving_variance:0')
conv1_eps = g.get_tensor_by_name('inference/guide/conv1/BatchNorm/batchnorm/add/y:0')
conv2_w = g.get_tensor_by_name('inference/guide/conv2/weights:0')
conv2_b = g.get_tensor_by_name('inference/guide/conv2/biases:0')
conv1w_, conv1b_, conv1mu_, conv1sigma_, conv1eps_, conv2w_, conv2b_ = sess.run(
[conv1_w, conv1_b, conv1_mu, conv1_sigma, conv1_eps, conv2_w, conv2_b])
conv1b_ -= conv1mu_/np.sqrt((conv1sigma_+conv1eps_))
conv1w_ = conv1w_/np.sqrt((conv1sigma_+conv1eps_))
conv1w_ = np.squeeze(conv1w_.astype(np.float32))
conv1b_ = np.squeeze(conv1b_.astype(np.float32))
conv1b_ = conv1b_[np.newaxis, :]
conv2w_ = np.squeeze(conv2w_.astype(np.float32))
conv2b_ = np.squeeze(conv2b_.astype(np.float32))
conv2 = np.append(conv2w_, conv2b_)
conv1 = np.vstack([conv1w_, conv1b_])
save(conv1.T, os.path.join(args.checkpoint_dir, 'guide_conv1.bin'))
save(conv2, os.path.join(args.checkpoint_dir, 'guide_conv2.bin'))
# Freeze the graph
input_graph_path = MODEL_NAME+'.pbtxt'
checkpoint_path = './'+MODEL_NAME+'.ckpt'
input_saver_def_path = ""
input_binary = False
output_node_names = "output/dense/BiasAdd"
restore_op_name = "save/restore_all"
filename_tensor_name = "save/Const:0"
output_frozen_graph_name = 'frozen_'+MODEL_NAME+'.pb'
output_optimized_graph_name = 'optimized_'+MODEL_NAME+'.pb'
clear_devices = True
freeze_graph.freeze_graph(input_graph_path, input_saver_def_path,
input_binary, checkpoint_path, output_node_names,
restore_op_name, filename_tensor_name,
output_frozen_graph_name, clear_devices, "")
# Optimize for inference
input_graph_def = tf.GraphDef()
with tf.gfile.Open(output_frozen_graph_name, "r") as f:
data = f.read()
input_graph_def.ParseFromString(data)
output_graph_def = optimize_for_inference_lib.optimize_for_inference(
input_graph_def,
["I"], # an array of the input node(s)
["O"], # an array of output nodes
def build(self, input_nodes=None, output_nodes=None):
if input_nodes is None:
input_nodes = self.gan.input_nodes()
if output_nodes is None:
output_nodes = self.gan.output_nodes()
save_file_text = self.name+".pbtxt"
build_file = os.path.expanduser("builds/"+save_file_text)
def create_path(filename):
return os.makedirs(os.path.expanduser(os.path.dirname(filename)), exist_ok=True)
create_path(build_file)
tf.train.write_graph(self.gan.session.graph, 'builds', save_file_text)
inputs = [x.name.split(":")[0] for x in input_nodes]
outputs = [x.name.split(":")[0] for x in output_nodes]
with self.gan.session as sess:
converter = tf.lite.TFLiteConverter.from_session(sess, self.gan.input_nodes(), self.gan.output_nodes())
tflite_model = converter.convert()
f = open("builds/"+ self.gan.name+".tflite", "wb")
f.write(tflite_model)
f.close()
tf.reset_default_graph()
self.gan.session.close()
[print("Input: ", x) for x in self.gan.input_nodes()]
[print("Output: ", y) for y in self.gan.output_nodes()]
print("Written to builds/"+self.gan.name+".tflite")
pbtxt_path = "builds/"+self.name +'.pbtxt'
checkpoint_path = "saves/"+self.name +'/model.ckpt'
input_saver_def_path = ""
input_binary = False
output_node_names = ",".join(outputs)
restore_op_name = "save/restore_all"
filename_tensor_name = "save/Const:0"
output_frozen_graph_name = 'builds/frozen_'+self.name +'.pb'
output_optimized_graph_name = 'builds/optimized_'+self.name+'.pb'
clear_devices = True
freeze_graph.freeze_graph(pbtxt_path, input_saver_def_path,
input_binary, checkpoint_path, output_node_names,
restore_op_name, filename_tensor_name,
output_frozen_graph_name, clear_devices, "")
input_graph_def = tf.GraphDef()
with tf.gfile.Open(output_frozen_graph_name, "rb") as f:
data = f.read()
input_graph_def.ParseFromString(data)
output_graph_def = optimize_for_inference_lib.optimize_for_inference(
input_graph_def,
inputs, # an array of the input node(s)
outputs, # an array of output nodes
tf.float32.as_datatype_enum)
# Save the optimized graph
f = tf.gfile.FastGFile(output_optimized_graph_name, "wb")
f.write(output_graph_def.SerializeToString())
f.flush()
f.close()
print("Saved generator to ", output_optimized_graph_name)
print("Testing loading ", output_optimized_graph_name)
with tf.gfile.FastGFile(output_optimized_graph_name, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
tf.import_graph_def(graph_def, name='')
#tflite_model = tf.lite.TFLiteConverter(graph_def, self.gan.input_nodes(), self.gan.output_nodes()).convert()
#f = open("builds/"+ self.gan.name+".tflite", "wb")
#f.write(tflite_model)
#f.close()
with tf.Session() as sess:
for input in inputs:
print("Input: ", input, sess.graph.get_tensor_by_name(input+":0"))
for output in outputs:
print("Output: ", output, sess.graph.get_tensor_by_name(output+":0"))
saver = tf.train.Saver()
saver.save(sess, save_path=checkpoint_path)
print('Checkpoint exported to {}'.format(checkpoint_path))
tf.train.write_graph(sess.graph_def, export_base_path, graph_name,
as_text=not as_binary)
print('Graph exported to {}'.format(graph_path))
if do_freeze:
print('Freezing graph...')
freeze_graph.freeze_graph(
input_graph=graph_path, input_saver='',
input_binary=as_binary, input_checkpoint=checkpoint_path,
output_node_names=output_node_name,
restore_op_name='save/restore_all',
filename_tensor_name='save/Const:0',
output_graph=frozen_graph_path, clear_devices=True,
initializer_nodes='')
print('Frozen graph exported to {}'.format(frozen_graph_path))
graph_path = frozen_graph_path
if do_optimize:
print('Optimizing graph...')
input_graph_def = tf.GraphDef()
with tf.gfile.Open(graph_path, 'rb') as f:
data = f.read()
input_graph_def.ParseFromString(data)
