def quantization(opt_graph, model_path, tf_records, eval_min_max_every_epoch):
# first_quantize
rewriter = quantize_graph.GraphRewriter(opt_graph, 'eightbit', None, None,
True)
first_quantize_graph = rewriter.rewrite(["policy_output", "value_output"])
if eval_min_max_every_epoch:
# insert_min_max_log
transform = 'insert_logging(op=RequantizationRange, show_name=true, message="__requant_min_max:")'
log_graph = TransformGraph(first_quantize_graph, ["pos_tensor"],
["policy_output", "value_output"],
[transform])
with tf.gfile.FastGFile(model_path + '_for_min_max.pb', 'wb') as f:
f.write(log_graph.SerializeToString())
# generate_min_max_log
with logged_timer('minmax time'):
generate_min_max_log(model_path + '_for_min_max.pb', tf_records,
model_path + 'log.txt')
# apply_calibration
transform = 'freeze_requantization_ranges(min_max_log_file="{0}")'.format(
model_path + 'log.txt')
calibration_graph = TransformGraph(first_quantize_graph, ["pos_tensor"],
["policy_output", "value_output"],
[transform])
# fuse_requantize
transform = 'fuse_quantized_conv_and_requantize strip_unused_nodes'
output_graph = TransformGraph(calibration_graph, ["pos_tensor"],
["policy_output", "value_output"],
[transform])
return output_graph
def prepare_for_dnn(sess,
graph_def,
in_node,
out_node,
out_graph,
dtype,
optimize=True,
quantize=False):
# Freeze graph. Replaces variables to constants.
graph_def = tf.graph_util.convert_variables_to_constants(
sess, graph_def, [out_node])
if optimize:
# Optimize graph. Removes training-only ops, unused nodes.
graph_def = optimize_for_inference_lib.optimize_for_inference(
graph_def, [in_node], [out_node], dtype.as_datatype_enum)
# Fuse constant operations.
transforms = ["fold_constants(ignore_errors=True)"]
if quantize:
transforms += ["quantize_weights(minimum_size=0)"]
transforms += ["sort_by_execution_order"]
graph_def = TransformGraph(graph_def, [in_node], [out_node],
transforms)
# Serialize
with tf.gfile.FastGFile(out_graph, 'wb') as f:
f.write(graph_def.SerializeToString())
def freeze_graph(input_checkpoint, output_graph):
output_node_names = "input_dataset/input_dataset_x,fc2/add" #输出的节点
input_node_names = 'input_dataset/input_dataset_x' #输入的节点
saver = tf.train.import_meta_graph(input_checkpoint + '.meta',
clear_devices=True)
graph = tf.get_default_graph() # 获得默认的图
input_graph_def = graph.as_graph_def() # 返回一个序列化的图代表当前的图
with tf.Session() as sess:
saver.restore(sess, input_checkpoint) # 恢复图并得到数据
# 输出所有Layer的tensor
# op = sess.graph.get_operations()
# [print(m.values()) for m in op][1]
output_graph_def = tf.graph_util.convert_variables_to_constants( # 模型持久化，将变量值固定
sess=sess,
input_graph_def=input_graph_def, # 等于:sess.graph_def
output_node_names=output_node_names.split(",")) # 如果有多个输出节点，以逗号隔开
# 压缩Graph，只保留输入输出
output_graph_def = optimize_for_inference_lib.optimize_for_inference(
output_graph_def, [input_node_names], output_node_names.split(","),
tf.float32.as_datatype_enum)
output_graph_def = TransformGraph(output_graph_def, [input_node_names],
output_node_names.split(","),
["sort_by_execution_order"])
with tf.gfile.GFile(output_graph, "wb") as f: # 保存模型
f.write(output_graph_def.SerializeToString()) # 序列化输出
def save(sess):
print(tf.all_variables())
input_graph_def = sess.graph.as_graph_def()
#for op in input_graph_def.node:
# print(op.name)
output_nodes_names=["init_26"]
output_graph_def = graph_util.convert_variables_to_constants(
sess, # The session
input_graph_def, # input_graph_def is useful for retrieving the nodes
output_nodes_names
)
output_graph_name="freeze.pb"
with tf.gfile.GFile(output_graph_name, "wb") as f:
f.write(output_graph_def.SerializeToString())
inp_node = ['Placeholder']
optimize_graph_def = optimize_for_inference_lib.optimize_for_inference(output_graph_def, [], output_nodes_names,
tf.float32.as_datatype_enum)
print("!")
optimize_graph_def = TransformGraph(optimize_graph_def, inp_node, output_nodes_names, ["sort_by_execution_order"])
output_graph_name="optimize.pb"
with tf.gfile.GFile(output_graph_name, "wb") as f:
f.write(optimize_graph_def.SerializeToString())
def save_as_pb(file_path,
sess,
input_node_names,
output_node_names,
as_text=False):
from google.protobuf import text_format
from tensorflow.python.framework import graph_util
from tensorflow.tools.graph_transforms import TransformGraph
graph_def = tf.get_default_graph().as_graph_def()
output_graph_def = graph_convert_variables_to_constants(
sess, graph_def, output_node_names)
transforms = [
"fold_constants(ignore_errors=true)", "fold_batch_norms",
"fold_old_batch_norms"
]
output_graph_def = TransformGraph(output_graph_def, input_node_names,
output_node_names, transforms)
with tf.gfile.GFile(file_path, "wb") as f:
if as_text:
f.write(text_format.MessageToString(output_graph_def))
else:
f.write(output_graph_def.SerializeToString())
return output_graph_def
def export_model(input_node_names, output_node_name):
freeze_graph.freeze_graph(OUTPUT_DIR + '/' + MODEL_NAME + '.pbtxt', None,
False, OUTPUT_DIR + '/' + MODEL_NAME + '.chkp',
output_node_name, "save/restore_all",
"save/Const:0",
OUTPUT_DIR + '/frozen_' + MODEL_NAME + '.pb',
True, "")
input_graph_def = tf.GraphDef()
with tf.gfile.Open(OUTPUT_DIR + '/frozen_' + MODEL_NAME + '.pb',
"rb") as f2:
input_graph_def.ParseFromString(f2.read())
transfroms = ['quantize_weights']
to_out = TransformGraph(input_graph_def, "input", output_node_name,
transfroms)
with tf.gfile.FastGFile(OUTPUT_DIR + '/trans_' + MODEL_NAME + '.pb',
"wb") as f3:
f3.write(to_out.SerializeToString())
print("transform succeed")
with tf.gfile.Open(OUTPUT_DIR + '/trans_' + 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(OUTPUT_DIR + '/opt_' + MODEL_NAME + '.pb',
"wb") as f:
f.write(output_graph_def.SerializeToString())
print("graph saved!")
def main(_):
args = parse_args()
saved_models = sorted(glob(os.path.join(args.saved_model_dir, '*')))
assert len(saved_models) > 0
with tf.Session() as sess:
tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], saved_models[-1])
graph = tf.get_default_graph()
convert_variables_to_constants = tf.graph_util.convert_variables_to_constants
output_graph_def = convert_variables_to_constants(
sess,
graph.as_graph_def(),
["transferred"])
input_names = ["images"]
output_names = ["transferred"]
transforms = ["strip_unused_nodes", "fold_batch_norms", "fold_constants", "quantize_weights"]
transformed_graph_def = TransformGraph(
output_graph_def,
input_names,
output_names,
transforms)
with tf.gfile.GFile(args.output_path, "wb") as f:
f.write(transformed_graph_def.SerializeToString())
def convert_to_pb(model,
path,
input_layer_name,
output_layer_name,
pbfilename,
verbose=False):
model.load(path, weights_only=True)
print("[INFO] Loaded CNN network weights from " + path + " ...")
print("[INFO] Re-export model ...")
del tf.get_collection_ref(tf.GraphKeys.TRAIN_OPS)[:]
model.save("model-tmp.tfl")
# taken from: https://stackoverflow.com/questions/34343259/is-there-an-example-on-how-to-generate-protobuf-files-holding-trained-tensorflow
print("[INFO] Re-import model ...")
input_checkpoint = "model-tmp.tfl"
saver = tf.train.import_meta_graph(input_checkpoint + '.meta', True)
sess = tf.Session()
saver.restore(sess, input_checkpoint)
# print out all layers to find name of output
if (verbose):
op = sess.graph.get_operations()
[print(m.values()) for m in op][1]
print("[INFO] Freeze model to " + pbfilename + " ...")
# freeze and removes nodes which are not related to feedforward prediction
minimal_graph = convert_variables_to_constants(sess,
sess.graph.as_graph_def(),
[output_layer_name])
graph_def = optimize_for_inference_lib.optimize_for_inference(
minimal_graph, [input_layer_name], [output_layer_name],
tf.float32.as_datatype_enum)
graph_def = TransformGraph(graph_def, [input_layer_name],
[output_layer_name],
["sort_by_execution_order"])
with tf.gfile.GFile(pbfilename, 'wb') as f:
f.write(graph_def.SerializeToString())
# write model to logs dir so we can visualize it as:
# tensorboard --logdir="logs"
if (verbose):
writer = tf.summary.FileWriter('logs', graph_def)
writer.close()
# tidy up tmp files
for f in glob.glob("model-tmp.tfl*"):
os.remove(f)
os.remove('checkpoint')
def load_graph(checkpoint_path, mb, seq_len):
init_all_op = tf.initialize_all_variables()
graph2 = tf.Graph()
with graph2.as_default():
with tf.Session(graph=graph2) as sess:
saver = tf.train.import_meta_graph(checkpoint_path + '.meta')
saver.restore(sess, checkpoint_path)
print("Restored structure...")
saver.restore(sess, checkpoint_path)
print("Restored params...")
'''
# input_names = ["IteratorGetNext"]
input_names = ["IteratorGetNext:0", "IteratorGetNext:1", "IteratorGetNext:4"]
output_names = ["loss/LogSoftmax"]
transforms = ['strip_unused_nodes(type=int32, shape="4,128")']
graph2 = TransformGraph(graph2.as_graph_def(), input_names, output_names, transforms)
# graph2 = TransformGraph(graph2, input_names, output_names, transforms)
# graph2 = tf.graph_util.remove_training_nodes(input_graph=graph2.as_graph_def())
graph2 = tf.graph_util.remove_training_nodes(input_graph=graph2)
'''
'''
input_names = ["IteratorGetNext:0", "IteratorGetNext:1", "IteratorGetNext:4"]
output_names = ["loss/LogSoftmax"]
transforms = ['strip_unused_nodes(type=int32, shape="4,128")']
# graph2 = TransformGraph(graph2.as_graph_def(), input_names, output_names, transforms)
graph2 = TransformGraph(graph2.as_graph_def(), inputs=input_names, outputs=output_names, transforms=transforms)
'''
'''
#2019-02-27 00:36:31.079753: I tensorflow/tools/graph_transforms/transform_graph.cc:317] Applying strip_unused_nodes
# terminate called after throwing an instance of 'std::out_of_range'
# what(): map::at
# Aborted
input_names = ["IteratorV2"] #Same result with "IteratorV2:0"
output_names = ["loss/LogSoftmax"]
transforms = ['strip_unused_nodes(type=resource)']
graph2 = TransformGraph(graph2.as_graph_def(), inputs=input_names, outputs=output_names, transforms=transforms)
'''
# input_names = ["IteratorGetNext", "IteratorGetNext:1", "IteratorGetNext:4"]
input_names = []
# output_names = ["loss/LogSoftmax"]
output_names = ["loss/Softmax"]
transforms = [
'strip_unused_nodes(type=int32, shape="' + str(mb) + ',' +
str(seq_len) + '")'
]
# graph2 = TransformGraph(graph2.as_graph_def(), input_names, output_names, transforms)
graph2 = TransformGraph(graph2.as_graph_def(),
inputs=input_names,
outputs=output_names,
transforms=transforms)
# for op in graph2.get_operations():
# print(op.name)
return graph2
def optimizePb(pbFilePath):
from tensorflow.python.tools import optimize_for_inference_lib
from tensorflow.tools.graph_transforms import TransformGraph
with tf.gfile.FastGFile(pbFilePath, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
graph_def = optimize_for_inference_lib.optimize_for_inference(graph_def, ['Placeholder'], ['final_result'], tf.float32.as_datatype_enum)
graph_def = TransformGraph(graph_def, ['module_apply_default/hub_input/Sub'], ['final_result'], ['remove_nodes(op=PlaceholderWithDefault)', 'strip_unused_nodes(type=float, shape=\"1,224,224,3\")', 'sort_by_execution_order'])
with tf.gfile.FastGFile('./inference_graph.pb', 'wb') as f:
f.write(graph_def.SerializeToString())
def optimize_model(model_path):
inputGraph = tf.GraphDef()
with tf.gfile.Open(model_path, 'rb') as model:
data2read = model.read()
inputGraph.ParseFromString(data2read)
with tf.variable_scope('lanenet'):
input_tensor = tf.placeholder(dtype=tf.float32, shape=[1, 256, 512, 3], name='input_tensor')
net = lanenet.LaneNet(phase='test', net_flag='vgg')
binary_seg_ret, instance_seg_ret = net.inference(input_tensor=input_tensor, name='lanenet_model')
with tf.variable_scope('lanenet/'):
binary_seg_ret = tf.cast(binary_seg_ret, dtype=tf.float32)
binary_seg_ret = tf.identity(binary_seg_ret, name= 'final_binary_output')
instance_seg_ret = tf.identity(instance_seg_ret, name='final_pixel_embedding_output')
# binary_seg_ret = tf.squeeze(binary_seg_ret, axis=0, name='final_binary_output') # 删除所有大小是 1 的维度
# instance_seg_ret = tf.squeeze(instance_seg_ret, axis=0, name='final_pixel_embedding_output')
"""outputGraph = optimize_for_inference_lib.optimize_for_inference(
inputGraph,
input_node_names=['lanenet/input_tensor'],
output_node_names=[
'lanenet/final_binary_output',
'lanenet/final_pixel_embedding_output'],
placeholder_type_enum=tf.int32.as_datatype_enum
)"""
outputGraph = TransformGraph(
inputGraph,
['lanenet/input_tensor'],
['lanenet/final_binary_output',
'lanenet/final_pixel_embedding_output'],
['remove_nodes(op=Identity, op=CheckNumerics)',
'merge_duplicate_nodes',
'strip_unused_nodes',
'fold_constants(ignore_errors=true)',
'fold_batch_norms',
'fold_old_batch_norms',
'quantize_weights',
'quantize_nodes',
'sort_by_execution_order']
)
new_name = model_path.split('/')[-2] + '/OptimizedGraph.pb'
model = tf.gfile.FastGFile(new_name, 'w')
model.write(outputGraph.SerializeToString())
def _quantize_static_tf_model(self, logdir, model_path, output_names):
with open(model_path, 'rb') as f:
serialized = f.read()
gdef = tf.GraphDef()
gdef.ParseFromString(serialized)
tf.reset_default_graph()
graph = tf.Graph()
with graph.as_default() as g:
transforms = [
"add_default_attributes",
"remove_nodes(op=Identity, op=CheckNumerics)",
"fold_constants(ignore_errors=true)", "fold_batch_norms",
"fold_old_batch_norms", "quantize_weights(minimum_size=1)",
"quantize_nodes", "strip_unused_nodes",
"sort_by_execution_order"
]
transformed_graph_def = TransformGraph(gdef, [], output_names,
transforms)
tf.import_graph_def(transformed_graph_def, name='')
tf.train.write_graph(graph,
logdir,
"./tf_quantized_frozen.pb",
as_text=False)
return os.path.join(logdir, 'tf_quantized_frozen.pb')
def save_graph_metadata(output_tensor, sess, feed_dict):
#First save the graph def
graph_def = tf.get_default_graph().as_graph_def()
transforms = [
'remove_nodes(op=Identity)',
'strip_unused_nodes',
'fold_batch_norms',
'fold_constants(ignore_errors=true)'
]
optimized_graph_def = TransformGraph(graph_def, [], [output_tensor.name], transforms)
with open('./graphDef.mtdata', 'w') as f:
f.write(str(optimized_graph_def))
# Save size information for tensors on which output depends
tensors_to_evaluate = []
tensors_to_evaluate_names = []
graph = tf.get_default_graph()
for node in optimized_graph_def.node:
cur_output = graph.get_operation_by_name(node.name).outputs[0]
tensors_to_evaluate.append(cur_output)
tensors_to_evaluate_names.append(node.name)
tensors_evaluated = sess.run(tensors_to_evaluate, feed_dict)
tensors_shape = list(map(lambda x : x.shape, tensors_evaluated))
# Write size info in a file
with open('./sizeInfo.mtdata','w') as f:
for ii, curr in enumerate(tensors_to_evaluate_names):
curShape = tensors_shape[ii]
f.write(tensors_to_evaluate_names[ii] + ' ')
for dim in curShape:
f.write(str(dim)+' ')
f.write('\n')
return optimized_graph_def
def ModelFreezer(model_fname, output_dir):
os.makedirs(output_dir, exist_ok=True)
if model_fname[-2:] == 'h5':
tf.keras.backend.set_learning_phase(0)
# need to define model if you have custom_object. model = get_model() then model.load_weights(H5_FILE). load_model just work if u don't have any custom_object
model = load_model(model_fname, compile=False, custom_objects={"swish": tf.nn.relu})
session = tf.keras.backend.get_session()
graph = session.graph.as_graph_def()
else:
graph = load_tf_graph(model_fname)
transformed_graph_def = TransformGraph(
graph,
inputs=[model.input.name[:-2]],
outputs=[model.output.name[:-2]],
transforms=_transform_ops()
)
const_graph_def = graph_util.convert_variables_to_constants(
session,
transformed_graph_def,
[model.output.name[:-2]]
)
try:
optimize_for_inference_lib.ensure_graph_is_valid(const_graph_def)
tf.train.write_graph(const_graph_def, output_dir, 'optimized_frozen.pb', as_text=False)
except ValueError as e:
print('Graph is invalid - {}'.format(e))
def _get_tensorflow_graph(model, num_outputs=1, quantize=False):
output_node_prefix = "output_node" #TODO check if this can be replaced with just using the Keras output name?
import tensorflow as tf
from keras import backend as K
K.set_learning_phase(0)
pred = [None] * num_outputs
pred_node_names = [None] * num_outputs
for i in range(num_outputs):
pred_node_names[i] = output_node_prefix + str(i)
pred[i] = tf.identity(model.outputs[i], name=pred_node_names[i])
print('output nodes names are: ', pred_node_names)
sess = K.get_session()
from tensorflow.python.framework import graph_util
if quantize:
from tensorflow.tools.graph_transforms import TransformGraph
transforms = ["quantize_weights", "quantize_nodes"]
transformed_graph_def = TransformGraph(sess.graph.as_graph_def(), [],
pred_node_names, transforms)
constant_graph = graph_util.convert_variables_to_constants(
sess, transformed_graph_def, pred_node_names)
else:
constant_graph = graph_util.convert_variables_to_constants(
sess, sess.graph.as_graph_def(), pred_node_names)
return constant_graph
def TacotronFreezer(checkpoint_dir, hp, output_dir):
tf.reset_default_graph()
synth = TacoSynthesizer()
checkpoint_path = tf.train.get_checkpoint_state(
checkpoint_dir).model_checkpoint_path
synth.load(checkpoint_path, hp)
os.makedirs(output_dir, exist_ok=True)
tf.train.write_graph(synth.session.graph,
output_dir,
'taco_variable_graph_def_ex.pd',
as_text=False)
transformed_graph_def = TransformGraph(
synth.session.graph.as_graph_def(),
inputs=_get_node_name([
synth.model.inputs, synth.model.input_lengths,
synth.model.split_infos
]),
outputs=_get_node_name(synth.model.tower_linear_outputs),
transforms=_transform_ops())
const_graph_def = graph_util.convert_variables_to_constants(
synth.session, transformed_graph_def,
_get_node_name(synth.model.tower_linear_outputs))
try:
optimize_for_inference_lib.ensure_graph_is_valid(const_graph_def)
tf.train.write_graph(const_graph_def,
output_dir,
'tacotron_optimized_frozen.pd',
as_text=False)
except Exception as e:
print('Graph is invalid: {}'.format(e))
def remove_dead_nodes(graph, in_list, out_list):
transforms = ['remove_nodes(op=Identity)', 'strip_unused_nodes']
optimized_graph_def = TransformGraph(graph.as_graph_def(), in_list,
out_list, transforms)
with tf.Graph().as_default() as opt_graph:
tf.import_graph_def(optimized_graph_def, name="")
return opt_graph
def strip_variable_init_constants(graph_def, input_tensor_names,
output_tensor_names):
transforms = [
"remove_nodes(op=Identity)",
"strip_unused_nodes",
]
# Sanity check if output/input nodes were constant and replaced with variables.
all_node_names = set([i.name for i in graph_def.node])
def get_true_names(tensor_names, all_nodes):
real_names = []
for i in tensor_names:
if i not in all_nodes:
var_name = i + "_mpc_const_var"
if var_name in all_nodes:
real_names.append(var_name)
else:
real_names.append(i)
return real_names
real_input_names = get_true_names(input_tensor_names, all_node_names)
real_output_names = get_true_names(output_tensor_names, all_node_names)
optimized_graph_def = TransformGraph(graph_def, real_input_names,
real_output_names, transforms)
return optimized_graph_def
def optimize_graph(self, transforms, optimized_file_name=None):
"""Applies graph transforms to the frozen GraphDef file & stores the optimized graph
Args
----
transforms(list)
Returns
-------
None
"""
if optimized_file_name is None:
file_name, extension = os.path.splitext(graph_filename)
self.optimized_file_name = file_name + "_optimized" + extension
else:
self.optimized_file_name = optimized_file_name
input_names = []
output_names = [self.output_node]
print("Loading TF GraphDef file")
graph_def = self.get_graph_def_from_file(
os.path.join(self.model_dir, self.graph_filename))
print("===============Graph Loaded==================")
print("===============Optimizing based on transforms : ")
print(transforms)
optimized_graph_def = TransformGraph(graph_def, input_names,
output_names, transforms)
tf.train.write_graph(
optimized_graph_def,
logdir=self.model_dir,
as_text=False,
name=optimized_file_name,
)
print("===============Graph Optimized==========")
def TacotronFreezer(checkpoint_path, hparams, output_dir):
tf.compat.v1.reset_default_graph()
synth = TacoSynthesizer()
synth.load(checkpoint_path, hparams)
os.makedirs(output_dir, exist_ok=True)
tf.io.write_graph(synth.session.graph, output_dir, 'taco_variable_graph_def.pb', as_text=False)
transformed_graph_def = TransformGraph(
synth.session.graph.as_graph_def(),
inputs=_get_node_name([synth.model.inputs, synth.model.input_lengths]),
outputs=_get_node_name([synth.mel_outputs]),
transforms=_transform_ops()
)
const_graph_def = graph_util.convert_variables_to_constants(
synth.session,
transformed_graph_def,
_get_node_name([synth.mel_outputs])
)
print('input_tensors: {}'.format(_get_node_name([synth.model.inputs, synth.model.input_lengths])))
print('output_tensors: {}'.format(_get_node_name([synth.mel_outputs])))
try:
optimize_for_inference_lib.ensure_graph_is_valid(const_graph_def)
tf.io.write_graph(const_graph_def, output_dir, 'optimized_frozen_tacotron.pb', as_text=False)
except ValueError as e:
print('Graph is invalid - {}'.format(e))
def main(quantize=False):
model = load_model(sys.argv[1])
model.summary()
orig_output_node_names = [node.op.name for node in model.outputs]
logging.info('Converted output node names are: {}'.format(
str(orig_output_node_names)))
sess = K.get_session()
saver = tf.train.Saver()
saver.save(sess, "./saved_ckpt/")
if quantize:
from tensorflow.tools.graph_transforms import TransformGraph
transforms = ['quantize_weights', 'quantize_nodes']
transform_graph_def = TransformGraph(sess.graph.as_graph_def(), [],
orig_output_node_names,
transforms)
else:
transform_graph_def = sess.graph.as_graph_def()
constant_graph = graph_util.convert_variables_to_constants(
sess, transform_graph_def, orig_output_node_names)
graph_io.write_graph(constant_graph,
"./saved_ckpt/",
"xor.pb",
as_text=False)
def save_model(fname, sess, graph=None):
def save(fname, graph_def):
pass
with tf.Graph().as_default() as g:
tf.import_graph_def(graph_def, name='')
graph_def = g.as_graph_def(add_shapes=True)
tf.train.write_graph(graph_def, ".", fname, as_text=False)
if graph == None:
graph_def = sess.graph_def
else:
graph_def = graph.as_graph_def(add_shapes=True)
input_nodes = [
'IteratorGetNext:0', 'IteratorGetNext:1', 'IteratorGetNext:2'
]
output_nodes = ['logits']
graph_def = graph_util.convert_variables_to_constants(
sess=sess, input_graph_def=graph_def, output_node_names=output_nodes)
graph_def = graph_util.remove_training_nodes(graph_def,
protected_nodes=output_nodes)
graph_def = optimize_for_inference_lib.optimize_for_inference(
graph_def, [], output_nodes, dtypes.float32.as_datatype_enum)
transforms = [
'remove_nodes(op=Identity, op=StopGradient)',
'fold_batch_norms',
'fold_old_batch_norms',
]
graph_def = TransformGraph(graph_def, input_nodes, output_nodes,
transforms)
save("build/data/bert_tf_v1_1_large_fp32_384_v2/model.pb", graph_def)
def main(args):
# If output_model path is relative and in cwd, make it absolute from root
output_model = FLAGS.output_model
if str(Path(output_model).parent) == '.':
output_model = str((Path.cwd() / output_model))
output_fld = Path(output_model).parent
output_model_name = Path(output_model).name
output_model_stem = Path(output_model).stem
output_model_pbtxt_name = output_model_stem + '.pbtxt'
# Create output directory if it does not exist
Path(output_model).parent.mkdir(parents=True, exist_ok=True)
if FLAGS.channels_first:
K.set_image_data_format('channels_first')
else:
K.set_image_data_format('channels_last')
model = load_model(FLAGS.input_model, FLAGS.input_model_json,
FLAGS.input_model_yaml)
orig_output_node_names = [node.op.name for node in model.outputs]
if FLAGS.output_nodes_prefix:
num_output = len(orig_output_node_names)
pred = [None] * num_output
converted_output_node_names = [None] * num_output
# Create dummy tf nodes to rename output
for i in range(num_output):
converted_output_node_names[i] = '{}{}'.format(
FLAGS.output_nodes_prefix, i)
pred[i] = tf.identity(model.outputs[i],
name=converted_output_node_names[i])
else:
converted_output_node_names = orig_output_node_names
logging.info('Converted output node names are: %s',
str(converted_output_node_names))
sess = K.get_session()
if FLAGS.output_meta_ckpt:
saver = tf.train.Saver()
saver.save(sess, str(output_fld / output_model_stem))
if FLAGS.save_graph_def:
tf.train.write_graph(sess.graph.as_graph_def(),
str(output_fld),
output_model_pbtxt_name,
as_text=True)
logging.info('Saved the graph definition in ascii format at %s',
str(Path(output_fld) / output_model_pbtxt_name))
if FLAGS.quantize:
from tensorflow.tools.graph_transforms import TransformGraph
transforms = ["quantize_weights", "quantize_nodes"]
transformed_graph_def = TransformGraph(sess.graph.as_graph_def(), [],
converted_output_node_names,
transforms)
constant_graph = graph_util.convert_variables_to_constants(
sess, transformed_graph_def, converted_output_node_names)
else:
constant_graph = graph_util.convert_variables_to_constants(
sess, sess.graph.as_graph_def(), converted_output_node_names)
graph_io.write_graph(constant_graph,
str(output_fld),
output_model_name,
as_text=False)
logging.info('Saved the frozen graph at %s',
str(Path(output_fld) / output_model_name))
def optimize_graph(model_file, output_dir, input_names = ['image_tensor'], output_names = ['num_detections', 'detection_classes', 'detection_scores', 'detection_boxes']):
print('Optimizing model {}...'.format(model_file))
graph_def = load_graph(model_file)
graph_stats(graph_def)
# TODO didn't test much of this operations
if LooseVersion(tf.__version__) >= LooseVersion('1.12.0'):
transforms = [
'strip_unused_nodes(type=float, shape="1,299,299,3")',
'remove_nodes(op=Identity, op=CheckNumerics)'
'fold_constants(ignore_errors=true)',
'fold_batch_norms',
'fold_old_batch_norms'
]
else:
print('[WARNING] Tensorflow version {} (< 1.12.0), some optimization disabled.'.format(tf.__version__))
transforms = [
'strip_unused_nodes(type=float, shape="1,299,299,3")',
'remove_nodes(op=CheckNumerics)'
'fold_constants(ignore_errors=true)',
'fold_batch_norms',
'fold_old_batch_norms'
]
graph_def_optimized = TransformGraph(graph_def, input_names, output_names, transforms)
print('\nAfter Optimization: ')
graph_stats(graph_def_optimized)
tf.train.write_graph(graph_def_optimized, logdir=output_dir, as_text=False, name='graph_optimized.pb')
def optimize_for_inference(frozen_graph_def: tf.GraphDef,
output_node_names: List,
graph_input: str) -> graph_pb2.GraphDef:
"""Optimize graph for inference.
Args:
frozen_graph_def: Frozen graph definition
output_node_names: Names of outputs
graph_input: Name of the image input to the graph.
Returns: Optimized inference graph definition.
"""
logging.info('Starting graph optimization.')
# Remove identity ops in initializers to allow fusing batch norm with conv in the next line
optimized_graph_def = tf.graph_util.remove_training_nodes(
frozen_graph_def)
optimized_graph_def = fold_batch_norms(optimized_graph_def)
transforms = [
'remove_nodes(op=Identity, op=CheckNumerics)',
'strip_unused_nodes', 'fold_constants(ignore_errors=true)'
]
optimized_graph_def = TransformGraph(optimized_graph_def,
[f"{graph_input}:0"],
output_node_names, transforms)
logging.info('Completed graph optimization.')
return optimized_graph_def
def freeze_graph_optimized():
input_checkpoint = "Neural_Network_Grades_simple2.ckpt"
output_node_names = [
"action/Relu",
]
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=True)
# We restore the weights
saver.restore(sess, input_checkpoint)
from tensorflow.tools.graph_transforms import TransformGraph
transforms = [
'add_default_attributes',
'remove_nodes(op=Identity, op=CheckNumerics)', 'fold_batch_norms',
'fold_old_batch_norms', 'strip_unused_nodes',
'sort_by_execution_order'
]
transformed_graph_def = TransformGraph(
tf.get_default_graph().as_graph_def(), 'state', output_node_names,
transforms)
# We use a built-in TF helper to export variables to constants
output_graph_def = tf.graph_util.convert_variables_to_constants(
sess, transformed_graph_def, output_node_names)
with tf.gfile.GFile("optimised_model.bytes", "wb") as f:
f.write(output_graph_def.SerializeToString())
def freeze_and_optimize_session(session,
keep_var_names=None,
input_names=None,
output_names=None,
clear_devices=True):
graph = session.graph
with graph.as_default():
freeze_var_names = list(
set(v.op.name for v in tf.compat.v1.global_variables()).difference(
keep_var_names or []))
output_names = output_names or []
output_names += [v.op.name for v in tf.compat.v1.global_variables()]
input_graph_def = graph.as_graph_def()
if clear_devices:
for node in input_graph_def.node:
node.device = ""
graph = tf.graph_util.remove_training_nodes(
input_graph_def, protected_nodes=output_names)
graph = tf.graph_util.convert_variables_to_constants(
session, graph, output_names, freeze_var_names)
transforms = [
'remove_nodes(op=Identity)',
'merge_duplicate_nodes',
'strip_unused_nodes',
'fold_constants(ignore_errors=true)',
'fold_batch_norms',
]
graph = TransformGraph(graph, input_names, output_names, transforms)
return graph
def tf_optimize(input_names, output_names, graph_def, fold_constant=True):
"""Extract inference subgraph and optimize graph."""
assert isinstance(input_names, list)
assert isinstance(output_names, list)
# TODO: is this needed ?
needed_names = [utils.node_name(i) for i in input_names] + \
[utils.node_name(i) for i in output_names]
graph_def = extract_sub_graph(graph_def, needed_names)
want_grappler = is_tf2() or LooseVersion(tf.__version__) >= "1.15"
if want_grappler:
graph_def = tf_optimize_grappler(input_names, output_names, graph_def,
fold_constant)
else:
# the older transform path
from tensorflow.tools.graph_transforms import TransformGraph # pylint: disable=redefined-outer-name
transforms = [
"fold_constants(ignore_errors=true)",
"remove_attribute(attribute_name=_class)", # remove node colocation attributes
"fold_batch_norms",
"fold_old_batch_norms",
]
graph_def = TransformGraph(graph_def, input_names, output_names,
transforms)
return graph_def
def _strip_unused_nodes(frozen_graph, concrete_func,
output_node_names):
# Find the names of the input nodes needed to extract the minimal
# inference graph. This is particularly useful for cases when the concrete
# function contains nodes that do not contribute the inference computation
# defined by the input/output pair. This would also eliminate op
# unsupported error caused by nodes outside of the minial infrerence
# graph.
input_node_names = []
for input_tensor in concrete_func.inputs:
if not input_tensor.dtype == 'resource':
op_name = input_tensor.name.split(':')[0]
# The graph freezing may turn the original inputs into constants, or
# remove them from the graph, so we need to ignore those.
try:
op = frozen_graph.get_operation_by_name(op_name)
if op.type != 'Const':
input_node_names.append(op_name)
except KeyError:
# The original input was removed when the graph was frozen.
continue
graph_transformations = ['strip_unused_nodes']
stripped_graph_def = TransformGraph(frozen_graph.as_graph_def(),
input_node_names,
output_node_names,
graph_transformations)
with tf.Graph().as_default() as stripped_graph:
tf.import_graph_def(stripped_graph_def, name='')
return stripped_graph
def _transform_graph(self, in_graph, out_graph, transforms):
"""Transforms input graph.
:param in_graph: input graph file or graphDef.
:param out_graph: output graph file or graphDef.
:param transforms: list of transforms.
:return:
"""
in_graph_def = in_graph if isinstance(
in_graph, tf.compat.v1.GraphDef) else self._read_graph(in_graph)
out_graph_def = TransformGraph(in_graph_def, self.inputs, self.outputs,
transforms)
if out_graph and not isinstance(out_graph, tf.compat.v1.GraphDef):
f = gfile.GFile(out_graph, 'wb')
f.write(out_graph_def.SerializeToString())
return out_graph_def
