def optimize_for_inference(input_graph_def, input_node_names, output_node_names,
placeholder_type_enum):
"""Applies a series of inference optimizations on the input graph.
Args:
input_graph_def: A GraphDef containing a training model.
input_node_names: A list of names of the nodes that are fed inputs during
inference.
output_node_names: A list of names of the nodes that produce the final
results.
placeholder_type_enum: The AttrValue enum for the placeholder data type, or
a list that specifies one value per input node name.
Returns:
An optimized version of the input graph.
"""
ensure_graph_is_valid(input_graph_def)
optimized_graph_def = input_graph_def
optimized_graph_def = strip_unused_lib.strip_unused(optimized_graph_def,
input_node_names,
output_node_names,
placeholder_type_enum)
optimized_graph_def = graph_util.remove_training_nodes(
optimized_graph_def, output_node_names)
optimized_graph_def = fold_batch_norms(optimized_graph_def)
optimized_graph_def = fuse_resize_and_conv(optimized_graph_def,
output_node_names)
ensure_graph_is_valid(optimized_graph_def)
return optimized_graph_def
def freeze(self):
gd = self.sess.graph.as_graph_def()
print("convt..")
for node in gd.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']
print("const...")
gd = graph_util.convert_variables_to_constants(self.sess, gd, self.outputNodes)
optlib.ensure_graph_is_valid(gd)
input_node_names = self.inputNodes
output_node_names = self.outputNodes
placeholder_type_enum = self.inputNodesTypes
for i in range(len(placeholder_type_enum)):
placeholder_type_enum[i] = placeholder_type_enum[i].as_datatype_enum
print("strip...")
gd = strip_unused_lib.strip_unused(gd, input_node_names, output_node_names, placeholder_type_enum)
optlib.ensure_graph_is_valid(gd)
filename = 'frozen ' + util.getTimeStamp() + '.pb'
tf.train.write_graph(gd, self.parentPath, filename, as_text=False)
return os.path.join(self.parentPath, filename)
def strip_subgraphs(self):
input_node_names = self._get_bare_input_names()
output_node_names = self._get_bare_output_names()
self.gdef = strip_unused_lib.strip_unused(
input_graph_def=self.original_gdef,
input_node_names=input_node_names,
output_node_names=output_node_names,
placeholder_type_enum=dtypes.float32.as_datatype_enum)
def strip_subgraphs(self) :
input_node_names = ['Preprocessor/sub']
output_node_names = ['concat', 'concat_1']
self.gdef = strip_unused_lib.strip_unused(
input_graph_def = self.original_gdef,
input_node_names = input_node_names,
output_node_names = output_node_names,
placeholder_type_enum = dtypes.float32.as_datatype_enum)
def optimize_graph(graph):
gdef = strip_unused_lib.strip_unused(
input_graph_def = graph.as_graph_def(),
input_node_names = [input_node],
output_node_names = [bbox_output_node, class_output_node],
placeholder_type_enum = dtypes.float32.as_datatype_enum)
with gfile.GFile(frozen_model_file, "wb") as f:
f.write(gdef.SerializeToString())
def __init__(self, frozen_file, inputshape, in_nodes, dest_nodes):
super(TensorflowParser2, self).__init__()
self.weight_loaded = True
# load model files into TensorFlow graph
with open(frozen_file, 'rb') as f:
serialized = f.read()
tensorflow.reset_default_graph()
original_gdef = tensorflow.GraphDef()
original_gdef.ParseFromString(serialized)
in_type_list = {}
for n in original_gdef.node:
if n.name in in_nodes:
in_type_list[n.name] = n.attr['dtype'].type
from tensorflow.python.tools import strip_unused_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.platform import gfile
original_gdef = strip_unused_lib.strip_unused(
input_graph_def = original_gdef,
input_node_names = in_nodes,
output_node_names = dest_nodes,
placeholder_type_enum = dtypes.float32.as_datatype_enum)
# Save it to an output file
frozen_model_file = './frozen.pb'
with gfile.GFile(frozen_model_file, "wb") as f:
f.write(original_gdef.SerializeToString())
with open(frozen_model_file, 'rb') as f:
serialized = f.read()
tensorflow.reset_default_graph()
model = tensorflow.GraphDef()
model.ParseFromString(serialized)
output_shape_map = dict()
input_shape_map = dict()
with tensorflow.Graph().as_default() as g:
input_map = {}
for i in range(len(inputshape)):
if in_type_list[in_nodes[i]] == 1:
dtype = tensorflow.float32
elif in_type_list[in_nodes[i]] == 3:
dtype = tensorflow.int32
x = tensorflow.placeholder(dtype, shape = [None] + inputshape[i])
input_map[in_nodes[i] + ':0'] = x
tensorflow.import_graph_def(model, name='', input_map=input_map)
with tensorflow.Session(graph = g) as sess:
meta_graph_def = tensorflow.train.export_meta_graph(filename='./my-model.meta')
model = meta_graph_def.graph_def
self.tf_graph = TensorflowGraph(model)
self.tf_graph.build()
def __init__(self,
meta_file,
checkpoint_file,
dest_nodes,
inputShape=None,
in_nodes=None):
super(TensorflowParser, self).__init__()
# load model files into TensorFlow graph
if meta_file:
model = TensorflowParser._load_meta(meta_file)
if checkpoint_file:
self.ckpt_data = TensorflowParser._load_weights(checkpoint_file)
self.weight_loaded = True
# extract subgraph using in_nodes and dest_nodes
if in_nodes != None and inputShape != None:
from tensorflow.python.tools import strip_unused_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.platform import gfile
input_node_names = in_nodes.split(',')
output_node_names = dest_nodes.split(',')
model = strip_unused_lib.strip_unused(
input_graph_def=model,
input_node_names=input_node_names,
output_node_names=output_node_names,
placeholder_type_enum=dtypes.float32.as_datatype_enum)
input_list = [None]
for i in range(len(inputShape)):
input_list.append(tensorflow.Dimension(inputShape[i]))
tensor_input = tensorflow.TensorShape(input_list)
# Build network graph
self.tf_graph = TensorflowGraph(model)
for node in self.tf_graph.model.node:
if node.name in input_node_names:
node.attr['shape'].list.shape.extend(
[tensor_input.as_proto()])
node.attr['_output_shapes'].list.shape.pop(
) #unknown_rank pop
node.attr['_output_shapes'].list.shape.extend(
[tensor_input.as_proto()])
# extract subgraph using dest_nodes
elif dest_nodes != None:
from tensorflow.python.framework.graph_util import extract_sub_graph
model = extract_sub_graph(model, dest_nodes.split(','))
self.tf_graph = TensorflowGraph(model)
else:
self.tf_graph = TensorflowGraph(model)
self.tf_graph.build()
def optimize_graph(input_path, output_path, input_nodes, output_nodes):
graph = tf.Graph()
with tf.Session(graph=graph) as sess:
tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], input_path)
gdef = strip_unused_lib.strip_unused(
input_graph_def=graph.as_graph_def(),
input_node_names=input_nodes,
output_node_names=output_nodes,
placeholder_type_enum=dtypes.float32.as_datatype_enum)
with gfile.GFile(output_path, 'wb') as f:
f.write(gdef.SerializeToString())
return graph
def freeze(self):
gd = self.sess.graph.as_graph_def()
print("convt..")
for node in gd.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']
print("const...")
gd = graph_util.convert_variables_to_constants(self.sess, gd,
["output"])
optlib.ensure_graph_is_valid(gd)
input_node_names = []
placeholder_type_enum = []
if (self.useLeft):
input_node_names.append("input_left")
placeholder_type_enum.append(tf.float32)
if (self.useRight):
input_node_names.append("input_left")
placeholder_type_enum.append(tf.float32)
if (self.useFace):
input_node_names.append("input_left")
placeholder_type_enum.append(tf.float32)
input_node_names.append("keep_prob")
placeholder_type_enum.append(tf.float32)
input_node_names.append("phase_train")
placeholder_type_enum.append(tf.bool)
output_node_names = ["output"]
for i in range(len(placeholder_type_enum)):
placeholder_type_enum[i] = placeholder_type_enum[
i].as_datatype_enum
print("strip...")
gd = strip_unused_lib.strip_unused(gd, input_node_names,
output_node_names,
placeholder_type_enum)
optlib.ensure_graph_is_valid(gd)
filename = 'frozen ' + time.strftime(R" %m-%d_%H-%M-%S",
time.localtime()) + '.pb'
tf.train.write_graph(gd, self.parentPath, filename, as_text=False)
return os.path.join(self.parentPath, filename)
def optimize_graph(input_path, output_path, input_nodes, output_nodes):
graph = tf.Graph()
with tf.Session(graph=graph) as sess:
tf.saved_model.loader.load(sess,
[tf.saved_model.tag_constants.SERVING],
input_path)
gdef = strip_unused_lib.strip_unused(
input_graph_def=graph.as_graph_def(),
input_node_names=input_nodes,
output_node_names=output_nodes,
placeholder_type_enum=dtypes.float32.as_datatype_enum)
with gfile.GFile(output_path, 'wb') as f:
f.write(gdef.SerializeToString())
return graph
def trim_graph_frozen(sess, graph_def, input_name_list, output_name_list, kill_norms=False):
# this is not robust
# Fix the batch-norm is_training issue
# TODO: clott
for n in graph_def.node:
if 'drop' in n.name:
print(n)
if kill_norms:
graph_def = convert_batchnorms(sess)
# Prune the unused nodes
graph_def = tf.graph_util.extract_sub_graph(graph_def, dest_nodes=output_name_list)
# for n in graph_def.node:
# if 'BatchNorm' in n.name and len(n.input) == 5 and not n.name.startswith(
# 'gradients') and 'Momentum' not in n.name:
# n.attr['is_training'].b = 0
# assert n.attr['is_training'].b == 0
# assert n.input[3].endswith('Const')
# assert n.input[4].endswith('Const_1')
# n.input[3] = n.input[3].rstrip('Const') + 'moving_mean/read'
# n.input[4] = n.input[4].rstrip('Const_1') + 'moving_variance/read'
gdef = strip_unused_lib.strip_unused(input_graph_def=graph_def,
input_node_names=input_name_list,
output_node_names=output_name_list,
placeholder_type_enum=dtypes.float32.as_datatype_enum)
gdef = tf.graph_util.convert_variables_to_constants(sess, gdef, output_name_list)
# output_graph = tf.GraphDef()
# output_graph.node.extend(graph_def.node)
#
# # Prune unnecessary placeholders & grad ops
# output_graph = tf.graph_util.convert_variables_to_constants(sess, output_graph, output_name_list)
# nodes = output_graph.node
# # Remove the dropout nodes
# src, dst = inspect(nodes)
# if src and dst:
# print(src, dst, nodes[src].name, nodes[dst].name)
# nodes[dst].input[0] = nodes[src].name
#
# nodes = prune(nodes, 'dropout')
# print('bring me back')
# # for i, n in enumerate(nodes):
# # if n.name == 'semantic':
# # n.input[0] = 'logits/semantic/BiasAdd'
# # Rebuild a graph-def
# graph_def = tf.GraphDef()
# graph_def.node.extend(nodes)
# graph_def = tf.graph_util.extract_sub_graph(graph_def, dest_nodes=output_name_list)
return gdef
def __init__(self, frozen_file, inputshape, in_nodes, dest_nodes):
super(TensorflowParser2, self).__init__()
self.weight_loaded = True
# load model files into TensorFlow graph
with open(frozen_file, 'rb') as f:
serialized = f.read()
tensorflow.reset_default_graph()
original_gdef = tensorflow.GraphDef()
original_gdef.ParseFromString(serialized)
# model = original_gdef
from tensorflow.python.tools import strip_unused_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.platform import gfile
input_node_names = in_nodes.split(',')
output_node_names = dest_nodes.split(',')
original_gdef = strip_unused_lib.strip_unused(
input_graph_def=original_gdef,
input_node_names=input_node_names,
output_node_names=output_node_names,
placeholder_type_enum=dtypes.float32.as_datatype_enum)
# Save it to an output file
frozen_model_file = './frozen.pb'
with gfile.GFile(frozen_model_file, "wb") as f:
f.write(original_gdef.SerializeToString())
with open(frozen_model_file, 'rb') as f:
serialized = f.read()
tensorflow.reset_default_graph()
model = tensorflow.GraphDef()
model.ParseFromString(serialized)
output_shape_map = dict()
input_shape_map = dict()
with tensorflow.Graph().as_default() as g:
x = tensorflow.placeholder(tensorflow.float32,
shape=[None] + inputshape)
tensorflow.import_graph_def(model,
name='',
input_map={in_nodes + ':0': x})
with tensorflow.Session(graph=g) as sess:
meta_graph_def = tensorflow.train.export_meta_graph(
filename='./my-model.meta')
model = meta_graph_def.graph_def
self.tf_graph = TensorflowGraph(model)
self.tf_graph.build()
def __init__(self,
meta_file,
checkpoint_file,
dest_nodes,
input_shape=None,
in_nodes=None,
input_format="NCHW".encode()):
graph_def = None
self.weights = None
self.inputs = in_nodes
self.outputs = dest_nodes
sess = tf.Session()
if meta_file is None:
raise Exception("meta_file must be provided")
new_saver = tf.train.import_meta_graph(meta_file)
if checkpoint_file is not None:
self.weights = dict()
new_saver.restore(sess,
tf.train.latest_checkpoint(checkpoint_file))
for var in tf.global_variables():
value = var.eval(sess)
self.weights[var.name.split(':')[0]] = value
self.infer = ModelInfer(sess)
graph_def, ver = tf.get_default_graph()._as_graph_def(add_shapes=True)
if in_nodes is not None and input_shape is not None:
graph_def = strip_unused_lib.strip_unused(
input_graph_def=graph_def,
input_node_names=in_nodes,
output_node_names=dest_nodes,
placeholder_type_enum=dtypes.float32.as_datatype_enum)
for node in graph_def.node:
if node.name in in_nodes:
index = in_nodes.index(node.name)
shape = [tf.Dimension(x) for x in input_shape[index]]
shape_proto = tf.TensorShape(shape).as_proto()
node.attr['_output_shapes'].list.shape.pop()
node.attr['_output_shapes'].list.shape.extend(
[shape_proto])
self.infer.gen_sample_data(node.name, input_shape[index])
self.tf_graph = TensorflowGraph(graph_def)
else:
raise Exception('in_nodes and output_nodes need be provided')
self.tf_graph.build(input_format)
def do_graph_freeze(output_file=None, output_node_names=None, variables_blacklist=''):
frozen = freeze_graph.freeze_graph_with_def_protos(
input_graph_def=tfv1.get_default_graph().as_graph_def(),
input_saver_def=saver.as_saver_def(),
input_checkpoint=checkpoint_path,
output_node_names=output_node_names,
restore_op_name=None,
filename_tensor_name=None,
output_graph=output_file,
clear_devices=False,
variable_names_blacklist=variables_blacklist,
initializer_nodes='')
input_node_names = []
return strip_unused_lib.strip_unused(
input_graph_def=frozen,
input_node_names=input_node_names,
output_node_names=output_node_names.split(','),
placeholder_type_enum=tf.float32.as_datatype_enum)
def optimize_for_inference(input_graph_def, input_node_names,
output_node_names, placeholder_type_enum):
"""Applies a series of inference optimizations on the input graph.
Args:
input_graph_def: A GraphDef containing a training model.
input_node_names: A list of names of the nodes that are fed inputs during
inference.
output_node_names: A list of names of the nodes that produce the final
results.
placeholder_type_enum: Data type of the placeholders used for inputs.
Returns:
An optimized version of the input graph.
"""
stripped_graph_def = strip_unused_lib.strip_unused(input_graph_def,
input_node_names,
output_node_names,
placeholder_type_enum)
detrained_graph_def = graph_util.remove_training_nodes(stripped_graph_def)
folded_graph_def = fold_batch_norms(detrained_graph_def)
return folded_graph_def
def optimize_for_inference(input_graph_def, input_node_names,
output_node_names, placeholder_type_enum):
"""Applies a series of inference optimizations on the input graph.
Args:
input_graph_def: A GraphDef containing a training model.
input_node_names: A list of names of the nodes that are fed inputs during
inference.
output_node_names: A list of names of the nodes that produce the final
results.
placeholder_type_enum: Data type of the placeholders used for inputs.
Returns:
An optimized version of the input graph.
"""
stripped_graph_def = strip_unused_lib.strip_unused(input_graph_def,
input_node_names,
output_node_names,
placeholder_type_enum)
detrained_graph_def = graph_util.remove_training_nodes(stripped_graph_def)
folded_graph_def = fold_batch_norms(detrained_graph_def)
return folded_graph_def
def strip_and_freeze_model(saved_model,
output_path,
input_node_names=[],
output_node_names=[]):
graph = tf.Graph()
with tf.Session(graph=graph) as sess:
print("loading model...")
tf.saved_model.loader.load(sess, [tf.saved_model.SERVING], saved_model)
print("stripping unused ops...")
gdef = strip_unused_lib.strip_unused(
input_graph_def=tf.get_default_graph().as_graph_def(),
input_node_names=input_node_names,
output_node_names=output_node_names,
placeholder_type_enum=dtypes.float32.as_datatype_enum,
)
gdef = tf.graph_util.convert_variables_to_constants(
sess, gdef, output_node_names)
with gfile.GFile(output_path, "wb") as f:
print("writing frozen model...")
f.write(gdef.SerializeToString())
return graph
with tf.compat.v1.gfile.GFile(mbv3_small_pb, "rb") as f:
graph_def = tf.compat.v1.GraphDef()
graph_def.ParseFromString(f.read())
with tf.Graph().as_default() as graph:
tf.import_graph_def(graph_def,
input_map=None,
return_elements=None,
name=""
)
operations = graph.get_operations()
input_graph_node = operations[2].outputs[0].name[:-2]
output_graph_node = operations[-1].outputs[0].name[:-2]
reduced_graph = strip_unused_lib.strip_unused(input_graph_def = graph.as_graph_def(),
input_node_names=[input_graph_node],
output_node_names=[output_graph_node],
placeholder_type_enum=dtypes.float32.as_datatype_enum)
with tf.compat.v1.gfile.GFile("mbv3_small_cleaned.pb", "wb") as f:
f.write(reduced_graph.SerializeToString())
#if we are converting from pb then we shouldn't strip the '/' from tensor names
mobilenetv3_small_mlmodel = tfcoreml.convert(
"mbv3_small_cleaned.pb",
output_feature_names = [output_graph_node],
input_name_shape_dict = {input_graph_node: [1, 224, 224, 3]},
image_input_names = [input_graph_node],
image_scale = 1/127.5,
red_bias = -1,
green_bias = -1,
blue_bias=-1,
def __init__(self, frozen_file, inputshape, in_nodes, dest_nodes):
if LooseVersion(tensorflow.__version__) < LooseVersion('1.8.0'):
raise ImportError(
'Your TensorFlow version %s is outdated. '
'MMdnn requires tensorflow>=1.8.0' % tensorflow.__version__)
super(TensorflowParser2, self).__init__()
self.weight_loaded = True
# load model files into TensorFlow graph
with open(frozen_file, 'rb') as f:
serialized = f.read()
tensorflow.reset_default_graph()
original_gdef = tensorflow.GraphDef()
original_gdef.ParseFromString(serialized)
in_type_list = {}
for n in original_gdef.node:
if n.name in in_nodes:
in_type_list[n.name] = n.attr['dtype'].type
from tensorflow.python.tools import strip_unused_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.platform import gfile
original_gdef = strip_unused_lib.strip_unused(
input_graph_def = original_gdef,
input_node_names = in_nodes,
output_node_names = dest_nodes,
placeholder_type_enum = dtypes.float32.as_datatype_enum)
# Save it to an output file
tempdir = tempfile.mkdtemp()
frozen_model_file = os.path.join(tempdir, 'frozen.pb')
with gfile.GFile(frozen_model_file, "wb") as f:
f.write(original_gdef.SerializeToString())
with open(frozen_model_file, 'rb') as f:
serialized = f.read()
shutil.rmtree(tempdir)
tensorflow.reset_default_graph()
model = tensorflow.GraphDef()
model.ParseFromString(serialized)
output_shape_map = dict()
input_shape_map = dict()
dtype = tensorflow.float32
with tensorflow.Graph().as_default() as g:
input_map = {}
for i in range(len(inputshape)):
dtype = TensorflowParser2.tf_dtype_map[in_type_list[in_nodes[i]]]
if in_type_list[in_nodes[i]] in (0, 1, 2):
x = tensorflow.placeholder(dtype, shape=[None] + inputshape[i])
elif in_type_list[in_nodes[i]] in (3, 4, 5, 6, 7):
x = tensorflow.placeholder(dtype, shape=inputshape[i])
elif in_type_list[in_nodes[i]] == 10:
x = tensorflow.placeholder(dtype)
else:
raise NotImplementedError
input_map[in_nodes[i] + ':0'] = x
tensorflow.import_graph_def(model, name='', input_map=input_map)
with tensorflow.Session(graph = g) as sess:
tempdir = tempfile.mkdtemp()
meta_graph_def = tensorflow.train.export_meta_graph(filename=os.path.join(tempdir, 'my-model.meta'))
model = meta_graph_def.graph_def
shutil.rmtree((tempdir))
self.tf_graph = TensorflowGraph(model)
self.tf_graph.build()
def __init__(self, frozen_file, inputshape, in_nodes, dest_nodes):
if tensorflow.__version__ < '1.8.0':
raise ImportError(
'Your TensorFlow version %s is outdated. '
'MMdnn requires tensorflow>=1.8.0' % tf.__version__)
super(TensorflowParser2, self).__init__()
self.weight_loaded = True
# load model files into TensorFlow graph
with open(frozen_file, 'rb') as f:
serialized = f.read()
tensorflow.reset_default_graph()
original_gdef = tensorflow.GraphDef()
original_gdef.ParseFromString(serialized)
in_type_list = {}
for n in original_gdef.node:
if n.name in in_nodes:
in_type_list[n.name] = n.attr['dtype'].type
from tensorflow.python.tools import strip_unused_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.platform import gfile
original_gdef = strip_unused_lib.strip_unused(
input_graph_def = original_gdef,
input_node_names = in_nodes,
output_node_names = dest_nodes,
placeholder_type_enum = dtypes.float32.as_datatype_enum)
# Save it to an output file
frozen_model_file = './frozen.pb'
with gfile.GFile(frozen_model_file, "wb") as f:
f.write(original_gdef.SerializeToString())
with open(frozen_model_file, 'rb') as f:
serialized = f.read()
tensorflow.reset_default_graph()
model = tensorflow.GraphDef()
model.ParseFromString(serialized)
output_shape_map = dict()
input_shape_map = dict()
with tensorflow.Graph().as_default() as g:
input_map = {}
for i in range(len(inputshape)):
if in_type_list[in_nodes[i]] == 1:
dtype = tensorflow.float32
elif in_type_list[in_nodes[i]] == 3:
dtype = tensorflow.int32
x = tensorflow.placeholder(dtype, shape = [None] + inputshape[i])
input_map[in_nodes[i] + ':0'] = x
tensorflow.import_graph_def(model, name='', input_map=input_map)
# graph_options = tensorflow.GraphOptions(
# optimizer_options=tensorflow.OptimizerOptions(
# opt_level=tensorflow.OptimizerOptions.L0, do_function_inlining=False))
# config = tensorflow.ConfigProto(graph_options=graph_options)
# with tensorflow.Session(graph = g, config=config) as sess:
with tensorflow.Session(graph = g) as sess:
meta_graph_def = tensorflow.train.export_meta_graph(filename='./my-model.meta')
model = meta_graph_def.graph_def
self.tf_graph = TensorflowGraph(model)
self.tf_graph.build()
def convert_to_core_ml(exported_graph_path, model_structure, output_path):
import tfcoreml
if not os.path.exists('.tmp'):
os.makedirs('.tmp')
if os.path.exists(output_path) and os.path.isdir(output_path):
shutil.rmtree(output_path)
os.makedirs(output_path)
if model_structure['type'] == ModelType.LOCALIZATION:
frozen_model_file = '.tmp/tmp_frozen_graph.pb'
with tf.Session(graph=tf.Graph()) as sess:
saved_model_path = os.path.join(exported_graph_path,
'saved_model/')
tf.saved_model.loader.load(sess,
[tf.saved_model.tag_constants.SERVING],
saved_model_path)
with gfile.GFile(frozen_model_file, 'wb') as f:
output_graph_def = tf.graph_util.convert_variables_to_constants(
sess,
tf.get_default_graph().as_graph_def(),
model_structure['output_names'])
input_node_names = ['Preprocessor/sub']
output_node_names = ['concat', 'concat_1']
stripped_graph_def = strip_unused_lib.strip_unused(
input_graph_def=output_graph_def,
input_node_names=input_node_names,
output_node_names=output_node_names,
placeholder_type_enum=dtypes.float32.as_datatype_enum)
f.write(stripped_graph_def.SerializeToString())
input_tensor_shapes = {'Preprocessor/sub:0': [1, 300, 300, 3]}
input_tensor_names = ['Preprocessor/sub:0']
output_tensor_names = ['concat:0', 'concat_1:0']
tfcoreml.convert(tf_model_path=frozen_model_file,
mlmodel_path=os.path.join(output_path,
'Model.mlmodel'),
output_feature_names=output_tensor_names,
red_bias=-1,
green_bias=-1,
blue_bias=-1,
image_scale=1.0 / 128.0,
input_name_shape_dict=input_tensor_shapes,
image_input_names=input_tensor_names)
json_labels = os.path.join(exported_graph_path, 'labels.json')
with open(json_labels) as f:
labels = json.load(f)
with open(os.path.join(output_path, 'Labels.swift'),
'w') as swift_labels:
swift_labels.write('//\n')
swift_labels.write('//\tLabels.swift\n')
swift_labels.write('//\tCloud Annotations\n')
swift_labels.write('//\n')
swift_labels.write('//\tGenerated on {}.\n'.format(
datetime.now().strftime("%m/%d/%y")))
swift_labels.write('//\n\n')
swift_labels.write('static let labels = [\n')
for label in labels:
swift_labels.write('\t"{}"\n'.format(label))
swift_labels.write(']\n')
anchors = os.path.join(exported_graph_path, 'Anchors.swift')
shutil.copy2(anchors, output_path)
else:
frozen_model_file = '.tmp/tmp_frozen_graph.pb'
with tf.Session(graph=tf.Graph()) as sess:
saved_model_path = os.path.join(exported_graph_path,
'saved_model/')
tf.saved_model.loader.load(sess,
[tf.saved_model.tag_constants.SERVING],
saved_model_path)
with gfile.GFile(frozen_model_file, 'wb') as f:
output_graph_def = tf.graph_util.convert_variables_to_constants(
sess,
tf.get_default_graph().as_graph_def(),
model_structure['output_names'])
f.write(output_graph_def.SerializeToString())
output_feature_names = [
'{}:0'.format(name) for name in model_structure['output_names']
]
class_labels = None
json_labels = os.path.join(exported_graph_path, 'labels.json')
text_labels = os.path.join(exported_graph_path, 'labels.txt')
if os.path.isfile(text_labels):
with open(text_labels, 'r') as f:
class_labels = f.read()
class_labels = list(
filter(bool,
[s.strip() for s in class_labels.splitlines()]))
elif os.path.isfile(json_labels):
with open(json_labels) as f:
class_labels = json.load(f)
tfcoreml.convert(
tf_model_path=frozen_model_file,
mlmodel_path=os.path.join(output_path, 'Model.mlmodel'),
output_feature_names=output_feature_names,
class_labels=class_labels,
red_bias=-1,
green_bias=-1,
blue_bias=-1,
image_scale=1.0 / 128.0,
image_input_names='{}:0'.format(model_structure['input_name']))
if __name__ == "__main__":
input_pb_file = '/opt/python-project/PycharmProjects/detect-test/models/face_mobile_8_18_512590_graph.pb'
output_pb_file = '/opt/python-project/PycharmProjects/detect-test/models/opt_512590.pb'
input_node_names = ["image_tensor"]
output_node_names = ['detection_boxes','detection_scores','detection_classes', 'num_detections']
input_graph_def = graph_pb2.GraphDef()
with open(input_pb_file, "rb") as f:
input_graph_def.ParseFromString(f.read())
output_graph_def = remove_assert_depend(input_graph_def)
output_graph_def = strip_unused_lib.strip_unused(
output_graph_def, input_node_names=input_node_names,
output_node_names=output_node_names,
placeholder_type_enum=dtypes.uint8.as_datatype_enum)
output_graph_def = remove_identity(output_graph_def, output_node_names)
skip_ops = []
new_ops = []
conv_names = ['FeatureExtractor/MobilenetV1/MobilenetV1/Conv2d_%s_pointwise/Conv2D'%i for i in range(1,14)] +\
['FeatureExtractor/MobilenetV1/MobilenetV1/Conv2d_%s_depthwise/depthwise'%i for i in range(1,14)]+\
["FeatureExtractor/MobilenetV1/MobilenetV1/Conv2d_0/Conv2D",
"FeatureExtractor/MobilenetV1/Conv2d_13_pointwise_1_Conv2d_2_1x1_256/Conv2D",
'FeatureExtractor/MobilenetV1/Conv2d_13_pointwise_2_Conv2d_2_3x3_s2_512/Conv2D',
'FeatureExtractor/MobilenetV1/Conv2d_13_pointwise_1_Conv2d_3_1x1_128/Conv2D',
'FeatureExtractor/MobilenetV1/Conv2d_13_pointwise_2_Conv2d_3_3x3_s2_256/Conv2D'
def __init__(self,
meta_file,
checkpoint_file,
dest_nodes,
inputShape=None,
in_nodes=None):
super(TensorflowParser, self).__init__()
# load model files into TensorFlow graph
if meta_file:
model = TensorflowParser._load_meta(meta_file)
if checkpoint_file:
self.ckpt_data = TensorflowParser._load_weights(checkpoint_file)
self.weight_loaded = True
# extract subgraph using in_nodes and dest_nodes
if in_nodes != None and inputShape != None:
from tensorflow.python.tools import strip_unused_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.platform import gfile
model = strip_unused_lib.strip_unused(
input_graph_def=model,
input_node_names=in_nodes,
output_node_names=dest_nodes,
placeholder_type_enum=dtypes.float32.as_datatype_enum)
input_list = [None]
for i in range(len(inputShape)):
input_list.append(tensorflow.Dimension(inputShape[i]))
tensor_input = tensorflow.TensorShape(input_list)
# Build network graph
self.tf_graph = TensorflowGraph(model)
for node in self.tf_graph.model.node:
if node.name in in_nodes:
node.attr['shape'].shape.CopyFrom(tensor_input.as_proto())
node.attr['_output_shapes'].list.shape.pop(
) #unknown_rank pop
node.attr['_output_shapes'].list.shape.extend(
[tensor_input.as_proto()])
# extract subgraph using dest_nodes
elif dest_nodes != None:
from tensorflow.python.framework.graph_util import extract_sub_graph
model = extract_sub_graph(model, dest_nodes)
self.tf_graph = TensorflowGraph(model)
else:
self.tf_graph = TensorflowGraph(model)
# Graph Transform
transforms = ["fold_constants(ignore_errors=true)"]
# Get input node name
if not in_nodes:
in_nodes = []
for node in model.node:
if node.op == 'Placeholder':
in_nodes.append(node.name)
transformed_graph_def = TransformGraph(model, in_nodes, dest_nodes,
transforms)
in_type_list = {}
in_shape_list = {}
for n in transformed_graph_def.node:
if n.name in in_nodes:
in_type_list[n.name] = n.attr['dtype'].type
in_node_shape = n.attr['shape'].shape
in_node_shape_str = self._shapeToStr(in_node_shape)
in_shape_list[n.name] = in_node_shape_str
dtype = tensorflow.float32
with tensorflow.Graph().as_default() as g:
input_map = {}
for in_node in in_nodes:
if in_type_list[in_node] == 1 or in_type_list[in_node] == 0:
dtype = tensorflow.float32
elif in_type_list[in_node] == 3:
dtype = tensorflow.int32
elif in_type_list[in_node] == 10:
dtype = tensorflow.bool
x = tensorflow.placeholder(dtype, shape=in_shape_list[in_node])
input_map[in_node] = x
tensorflow.import_graph_def(transformed_graph_def,
name='',
input_map=input_map)
with tensorflow.Session(graph=g) as sess:
tempdir = tempfile.mkdtemp()
meta_graph_def = tensorflow.train.export_meta_graph(
filename=os.path.join(tempdir, 'my-model.meta'))
model = meta_graph_def.graph_def
shutil.rmtree(tempdir)
self.tf_graph = TensorflowGraph(model)
self.tf_graph.build()
process_graph(self.tf_graph, self.ckpt_data)
def __init__(self, frozen_file, inputshape, in_nodes, dest_nodes):
super(TensorflowParser2, self).__init__()
self.weight_loaded = True
# load model files into TensorFlow graph
with open(frozen_file, 'rb') as f:
serialized = f.read()
tensorflow.reset_default_graph()
original_gdef = tensorflow.GraphDef()
original_gdef.ParseFromString(serialized)
# model = original_gdef
from tensorflow.python.tools import strip_unused_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.platform import gfile
input_node_names = in_nodes.split(',')
output_node_names = dest_nodes.split(',')
gdef = strip_unused_lib.strip_unused(
input_graph_def=original_gdef,
input_node_names=input_node_names,
output_node_names=output_node_names,
placeholder_type_enum=dtypes.float32.as_datatype_enum)
# Save it to an output file
frozen_model_file = './frozen.pb'
with gfile.GFile(frozen_model_file, "wb") as f:
f.write(gdef.SerializeToString())
with open(frozen_model_file, 'rb') as f:
serialized = f.read()
tensorflow.reset_default_graph()
model = tensorflow.GraphDef()
model.ParseFromString(serialized)
output_shape_map = dict()
input_shape_map = dict()
with tensorflow.Graph().as_default() as g:
tensorflow.import_graph_def(model, name='')
ops = g.get_operations()
N = len(ops)
p = 0
for i in range(N):
for x in ops[i].inputs:
input_shape_map[x.name] = x.get_shape()
for x in ops[i].outputs:
output_shape_map[x.name] = x.get_shape()
# for pytest
if type(inputshape[2]) == int:
tensor_input = tensorflow.TensorShape([
None,
tensorflow.Dimension(inputshape[0]),
tensorflow.Dimension(inputshape[1]),
tensorflow.Dimension(inputshape[2])
])
output_shape_map[input_node_names[0]] = tensor_input
else:
if len(inputshape[2].split(',')) == 2:
param = inputshape[2].split(',')[1]
tensor_input_2 = tensorflow.TensorShape(
[tensorflow.Dimension(param)])
output_shape_map[input_node_names[1]] = tensor_input_2
tensor_input = tensorflow.TensorShape([
None,
tensorflow.Dimension(inputshape[0]),
tensorflow.Dimension(inputshape[1]),
tensorflow.Dimension(inputshape[2].split(',')[0])
])
output_shape_map[input_node_names[0]] = tensor_input
else:
tensor_input = tensorflow.TensorShape([
None,
tensorflow.Dimension(inputshape[0]),
tensorflow.Dimension(inputshape[1]),
tensorflow.Dimension(inputshape[2])
])
output_shape_map[input_node_names[0]] = tensor_input
# assert False
self.tf_graph = TensorflowGraph(model)
for node in self.tf_graph.model.node:
if (node.name +
':0') in output_shape_map and node.op != 'Placeholder':
node.attr['_output_shapes'].list.shape.extend(
[output_shape_map[node.name + ':0'].as_proto()])
if node.op == 'MirrorPad':
node.attr['paddings'].list.shape.extend(
[input_shape_map[node.name + '/paddings:0'].as_proto()])
if node.op == 'QuantizeV2':
node.attr['shape'].list.shape.extend(
[input_shape_map[node.name + ':0'].as_proto()])
if node.op == 'RequantizationRange':
map_key = node.name.split(
'eightbit')[0] + "eightbit_quantized_conv:0"
node.attr['shape'].list.shape.extend(
[input_shape_map[map_key].as_proto()])
if node.op == 'Requantize':
map_key = node.name.replace("requantize",
"quantized_conv") + ":0"
node.attr['shape'].list.shape.extend(
[input_shape_map[map_key].as_proto()])
if node.name in input_node_names:
if node.name in output_shape_map.keys():
node.attr['shape'].list.shape.extend(
[output_shape_map[node.name].as_proto()])
node.attr['_output_shapes'].list.shape.extend(
[output_shape_map[node.name].as_proto()])
self.tf_graph.build()
def __init__(self,
pb_file,
dest_nodes,
input_shape=None,
in_nodes=None,
input_format="NCHW".encode()):
with open(pb_file, 'rb') as f:
serialized = f.read()
tf.reset_default_graph()
original_graph_def = tf.GraphDef()
original_graph_def.ParseFromString(serialized)
self.inputs = list()
self.outputs = dest_nodes
sess = tf.Session(graph=tf.get_default_graph())
sess.run(tf.global_variables_initializer())
self.infer = ModelInfer(sess)
original_graph_def = strip_unused_lib.strip_unused(
input_graph_def=original_graph_def,
input_node_names=in_nodes,
output_node_names=dest_nodes,
placeholder_type_enum=dtypes.float32.as_datatype_enum)
graph_def = tf.GraphDef()
graph_def.ParseFromString(original_graph_def.SerializeToString())
in_type_list = dict()
for node in graph_def.node:
if node.name in in_nodes:
in_type_list[node.name] = node.attr['dtype'].type
input_shape = list(input_shape)
if not isinstance(input_shape[0], list):
input_shape = [input_shape]
input_map = dict()
for i in range(len(input_shape)):
if in_type_list[in_nodes[i]] == 1 or in_type_list[
in_nodes[i]] == 0:
dtype = tf.float32
x = tf.placeholder(dtype, shape=input_shape[i])
elif in_type_list[in_nodes[i]] == 3:
dtype = tf.int32
x = tf.placehoder(dtype, shape=input_shape[i])
else:
raise Exception("Unexpected dtype for input, only support " \
"float32 and int32 now")
input_map[in_nodes[i] + ":0"] = x
self.inputs.append(x.name.split(':')[0])
self.infer.gen_sample_data(x.name, input_shape[i])
tf.import_graph_def(graph_def, name="", input_map=input_map)
graph_def = tf.get_default_graph()._as_graph_def(add_shapes=True)[0]
self.tf_graph = TensorflowGraph(graph_def)
self.tf_graph.build(input_format)
self.weights = dict()
for node in graph_def.node:
if node.op.lower() == "const":
try:
node.attr['value'].tensor.tensor_content
weight = tensor_util.MakeNdarray(node.attr['value'].tensor)
self.weights[node.name] = weight
except:
continue
tf.reset_default_graph()
original_gdef = tf.GraphDef()
original_gdef.ParseFromString(serialized)
with tf.Graph().as_default() as g:
tf.import_graph_def(original_gdef, name='')
# Strip unused subgraphs and save it as another frozen TF model
from tensorflow.python.tools import strip_unused_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.platform import gfile
input_node_names = ['Preprocessor/sub']
output_node_names = ['concat', 'concat_1']
gdef = strip_unused_lib.strip_unused(
input_graph_def=original_gdef,
input_node_names=input_node_names,
output_node_names=output_node_names,
placeholder_type_enum=dtypes.float32.as_datatype_enum)
# Save the feature extractor to an output file
frozen_model_file = 'core_ml_conversion/ssd_mobilenet_feature_extractor.pb'
with gfile.GFile(frozen_model_file, "wb") as f:
f.write(gdef.SerializeToString())
# Now we have a TF model ready to be converted to CoreML
import tfcoreml
# Supply a dictionary of input tensors' name and shape (with # batch axis)
input_tensor_shapes = {
"Preprocessor/sub:0": [1, 300, 300, 3]
} # batch size is 1
# Output CoreML model path
coreml_model_file = 'ssd_mobilenet.mlmodel'
def __init__(self, frozen_file, inputshape, in_nodes, dest_nodes):
super(TensorflowParser2, self).__init__()
self.weight_loaded = True
# load model files into TensorFlow graph
with open(frozen_file, 'rb') as f:
serialized = f.read()
tensorflow.reset_default_graph()
original_gdef = tensorflow.GraphDef()
original_gdef.ParseFromString(serialized)
# model = original_gdef
from tensorflow.python.tools import strip_unused_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.platform import gfile
input_node_names = in_nodes.split(',')
output_node_names = dest_nodes.split(',')
gdef = strip_unused_lib.strip_unused(
input_graph_def = original_gdef,
input_node_names = input_node_names,
output_node_names = output_node_names,
placeholder_type_enum = dtypes.float32.as_datatype_enum)
# Save it to an output file
frozen_model_file = './frozen.pb'
with gfile.GFile(frozen_model_file, "wb") as f:
f.write(gdef.SerializeToString())
with open(frozen_model_file, 'rb') as f:
serialized = f.read()
tensorflow.reset_default_graph()
model = tensorflow.GraphDef()
model.ParseFromString(serialized)
output_shape_map = dict()
input_shape_map = dict()
with tensorflow.Graph().as_default() as g:
tensorflow.import_graph_def(original_gdef, name='')
ops = g.get_operations()
N = len(ops)
p = 0
for i in range(N):
for x in ops[i].inputs:
input_shape_map[x.name] = x.get_shape()
for x in ops[i].outputs:
output_shape_map[x.name] = x.get_shape()
tensor_input = tensorflow.TensorShape([None, tensorflow.Dimension(inputshape[0]), tensorflow.Dimension(inputshape[1]), tensorflow.Dimension(inputshape[2])])
output_shape_map[input_node_names[0]] = tensor_input
self.tf_graph = TensorflowGraph(model)
for node in self.tf_graph.model.node:
if node.name == 'input' and node.op == 'Placeholder':
node.attr['shape'].list.shape.extend([output_shape_map[node.name + ':0'].as_proto()])
if (node.name + ':0') in output_shape_map:
node.attr['_output_shapes'].list.shape.extend([output_shape_map[node.name + ':0'].as_proto()])
if node.op == 'MirrorPad':
node.attr['paddings'].list.shape.extend([input_shape_map[node.name + '/paddings:0'].as_proto()])
if node.op == 'QuantizeV2':
node.attr['shape'].list.shape.extend([input_shape_map[node.name + ':0'].as_proto()])
if node.op == 'RequantizationRange':
map_key = node.name.split('eightbit')[0] + "eightbit_quantized_conv:0"
node.attr['shape'].list.shape.extend([input_shape_map[map_key].as_proto()])
if node.op == 'Requantize':
map_key = node.name.replace("requantize", "quantized_conv")+":0"
node.attr['shape'].list.shape.extend([input_shape_map[map_key].as_proto()])
if node.name == in_nodes:
node.attr['shape'].list.shape.extend([output_shape_map[node.name].as_proto()])
self.tf_graph.build()