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 tf_optimize(sess, inputs, outputs, graph_def):
"""Optimize tensorflow graph for inference."""
transforms = [
"fold_constants(ignore_errors=true)",
"fold_batch_norms",
"fold_old_batch_norms",
]
needed_names = [utils.node_name(i) for i in inputs] + [utils.node_name(i) for i in outputs]
graph_def = graph_util.extract_sub_graph(graph_def, needed_names)
graph_def = TransformGraph(graph_def, inputs, outputs, transforms)
return graph_def
def tf_optimize(inputs, outputs, graph_def, fold_constant=None):
"""Optimize tensorflow graph for inference."""
transforms = []
if fold_constant:
transforms.extend([
"fold_constants(ignore_errors=true)",
"remove_attribute(attribute_name=_class)", # remove node colocation attributes
])
transforms.extend([
"fold_batch_norms",
"fold_old_batch_norms",
])
needed_names = [utils.node_name(i)
for i in inputs] + [utils.node_name(i) for i in outputs]
graph_def = graph_util.extract_sub_graph(graph_def, needed_names)
graph_def = TransformGraph(graph_def, inputs, outputs, transforms)
return graph_def
def remove_redundant_inputs(frozen_graph, input_names):
"""Remove redundant inputs not in frozen graph."""
frozen_inputs = []
# get inputs in frozen graph
node_names = set(n.name for n in frozen_graph.node)
frozen_inputs = [inp for inp in input_names if utils.node_name(inp) in node_names]
deleted_inputs = list(set(input_names) - set(frozen_inputs))
if deleted_inputs:
logger.warning("inputs [%s] is not in frozen graph, delete them", ",".join(deleted_inputs))
return frozen_inputs
def tf_optimize(sess, inputs, outputs, graph_def):
# print("tf_optimize begin")
"""Optimize tensorflow graph for inference."""
transforms = [
"fold_constants(ignore_errors=true)",
"fold_batch_norms",
"fold_old_batch_norms",
]
# TODO 这俩 在 研究 研究
needed_names = [utils.node_name(i) for i in inputs] + [utils.node_name(i) for i in outputs]
print("---------------needed_names:", needed_names)
graph_def = graph_util.extract_sub_graph(graph_def, needed_names)
print("extract_sub_graph done")
graph_def = TransformGraph(graph_def, inputs, outputs, transforms)
print("TransformGraph done")
return graph_def
def remove_redundant_inputs(frozen_graph, input_names):
"""Remove redundant inputs not in frozen graph."""
frozen_inputs = []
# get inputs in frozen graph
for n in frozen_graph.node:
for inp in input_names:
if utils.node_name(inp) == n.name:
frozen_inputs.append(inp)
deleted_inputs = list(set(input_names) - set(frozen_inputs))
if deleted_inputs:
logger.warning("inputs [%s] is not in frozen graph, delete them", ",".join(deleted_inputs))
return frozen_inputs
def tf_optimize(input_tensors, output_tensors, graph_def, fold_constant=True):
"""Extract inference subgraph and optimize graph."""
assert isinstance(input_tensors, dict)
assert isinstance(output_tensors, dict)
try:
input_tensors = {
name: tensor for name, tensor in input_tensors.items()
if tensor.dtype != tf.dtypes.resource
}
except: # pylint: disable=bare-except
pass
# TODO: is this needed ?
needed_names = [utils.node_name(i) for i in input_tensors.keys()] + \
[utils.node_name(i) for i in output_tensors.keys()]
graph_def = extract_sub_graph(graph_def, needed_names)
if fold_constant:
want_grappler = is_tf2() or LooseVersion(tf.__version__) >= "1.15"
if want_grappler:
graph_def = tf_optimize_grappler(input_tensors, output_tensors, graph_def, fold_constant)
else:
# the older transform path
from tensorflow.tools.graph_transforms import TransformGraph # pylint: disable=redefined-outer-name
transforms = []
if fold_constant:
transforms.extend([
"fold_constants(ignore_errors=true)",
"remove_attribute(attribute_name=_class)", # remove node colocation attributes
])
transforms.extend([
"fold_batch_norms",
"fold_old_batch_norms",
])
graph_def = TransformGraph(graph_def, input_tensors.keys(), output_tensors.keys(), transforms)
return graph_def
def get_node_by_name(self, name):
"""Get node by name."""
ret = self._nodes_by_name.get(name)
if not ret:
ret = self._nodes_by_name.get(node_name(name))
if not ret:
# if we processed the graph fully, set_nodes() the graph has no longer const nodes
# since we moved them to be initializers. But all graph processing code uses Node
# as the common data structure. To avoid special casing lots of code for initializers
# we create a dummy 'Const' Node here.
initializer = self._initializers.get(name)
if initializer:
ret = Node(helper.make_node("Const", [], [name], name=name, value=initializer),
self, skip_conversion=True)
return ret
def from_saved_model(model_path, input_names, output_names):
"""Load tensorflow graph from saved_model."""
# make sure we start with clean default graph
tf.reset_default_graph()
inputs = {}
outputs = {}
try:
from tensorflow.contrib.saved_model.python.saved_model import signature_def_utils
# pylint: disable=unnecessary-lambda
get_signature_def = lambda meta_graph_def, k: \
signature_def_utils.get_signature_def_by_key(meta_graph_def, k)
except ImportError:
# TF1.12 changed the api
get_signature_def = lambda meta_graph_def, k: meta_graph_def.signature_def[
k]
with tf.Session() as sess:
meta_graph_def = tf.saved_model.loader.load(
sess, [tf.saved_model.tag_constants.SERVING], model_path)
for k in meta_graph_def.signature_def.keys():
inputs_tensor_info = get_signature_def(meta_graph_def, k).inputs
for _, input_tensor in sorted(inputs_tensor_info.items()):
inputs[input_tensor.name] = sess.graph.get_tensor_by_name(
input_tensor.name)
outputs_tensor_info = get_signature_def(meta_graph_def, k).outputs
for _, output_tensor in sorted(outputs_tensor_info.items()):
outputs[output_tensor.name] = sess.graph.get_tensor_by_name(
output_tensor.name)
frozen_graph = freeze_session(sess, output_names=list(outputs.keys()))
frozen_inputs = []
# get inputs in frozen graph
for n in frozen_graph.node:
for inp, _ in inputs.items():
if utils.node_name(inp) == n.name:
frozen_inputs.append(inp)
deleted_inputs = list(set(inputs.keys()) - set(frozen_inputs))
if deleted_inputs:
log.warning("inputs [%s] is not in frozen graph, delete them",
",".join(deleted_inputs))
# clean up
tf.reset_default_graph()
return frozen_graph, frozen_inputs, outputs.keys()
def run_test_case(self,
func,
feed_dict,
input_names_with_port,
output_names_with_port,
rtol=1e-07,
atol=1e-5,
convert_var_to_const=True,
constant_fold=True,
check_value=True,
check_shape=True,
check_dtype=True,
process_args=None,
onnx_feed_dict=None,
graph_validator=None,
as_session=False,
large_model=False):
# optional - passed to process_tf_graph
if process_args is None:
process_args = {}
# optional - pass distinct feed_dict to onnx runtime
if onnx_feed_dict is None:
onnx_feed_dict = feed_dict
input_names_with_port = list(feed_dict)
tf_reset_default_graph()
graph_def = None
np.random.seed(1) # Make it reproducible.
clean_feed_dict = {utils.node_name(k): v for k, v in feed_dict.items()}
if is_tf2() and not as_session:
#
# use eager to execute the tensorflow func
#
# numpy doesn't work for all ops, make it tf.Tensor()
input_tensors = [
tf.TensorSpec(shape=v.shape,
dtype=tf.as_dtype(v.dtype),
name=utils.node_name(k))
for k, v in feed_dict.items()
]
input_list = [
tf.convert_to_tensor(v,
dtype=tf.as_dtype(v.dtype),
name=utils.node_name(k))
for k, v in feed_dict.items()
]
tf.random.set_seed(1)
expected = func(*input_list)
if isinstance(expected, (list, tuple)):
# list or tuple
expected = [x.numpy() for x in expected]
else:
# single result
expected = [expected.numpy()]
# now make the eager functions a graph
concrete_func = tf.function(func,
input_signature=tuple(input_tensors))
concrete_func = concrete_func.get_concrete_function()
graph_def = from_function(concrete_func,
input_names=list(feed_dict.keys()),
output_names=output_names_with_port,
large_model=large_model)
else:
#
# use graph to execute the tensorflow func
#
with tf_session() as sess:
tf_set_random_seed(1)
input_list = []
for k, v in clean_feed_dict.items():
input_list.append(
tf_placeholder(name=k,
shape=v.shape,
dtype=tf.as_dtype(v.dtype)))
func(*input_list)
variables_lib.global_variables_initializer().run()
tf_tables_initializer().run()
output_dict = []
for out_name in output_names_with_port:
output_dict.append(sess.graph.get_tensor_by_name(out_name))
expected = sess.run(output_dict, feed_dict=feed_dict)
graph_def = freeze_session(sess,
input_names=list(feed_dict.keys()),
output_names=output_names_with_port)
tf_reset_default_graph()
with tf_session() as sess:
tf.import_graph_def(graph_def, name='')
graph_def = tf_optimize(list(feed_dict.keys()),
output_names_with_port,
graph_def,
fold_constant=constant_fold)
tf_reset_default_graph()
with tf_session() as sess:
const_node_values = None
if large_model:
const_node_values = compress_graph_def(graph_def)
tf.import_graph_def(graph_def, name='')
if self.config.is_debug_mode:
model_path = os.path.join(
self.test_data_directory,
self._testMethodName + "_after_tf_optimize.pb")
utils.save_protobuf(model_path, graph_def)
self.logger.debug("created file %s", model_path)
g = process_tf_graph(sess.graph,
opset=self.config.opset,
input_names=list(feed_dict.keys()),
output_names=output_names_with_port,
target=self.config.target,
const_node_values=const_node_values,
**process_args)
g = optimizer.optimize_graph(g)
actual = self.run_backend(g, output_names_with_port,
onnx_feed_dict, large_model)
for expected_val, actual_val in zip(expected, actual):
if check_value:
self.assertAllClose(expected_val,
actual_val,
rtol=rtol,
atol=atol)
if check_dtype:
self.assertEqual(expected_val.dtype, actual_val.dtype)
# why need shape checke: issue when compare [] with scalar
# https://github.com/numpy/numpy/issues/11071
if check_shape:
self.assertEqual(expected_val.shape, actual_val.shape)
if graph_validator:
self.assertTrue(graph_validator(g))
return g
def freeze_and_run_tf(self, func, feed_dict, outputs, as_session,
premade_placeholders, large_model, constant_fold):
np.random.seed(1) # Make it reproducible.
clean_feed_dict = {utils.node_name(k): v for k, v in feed_dict.items()}
if is_tf2() and not as_session:
#
# use eager to execute the tensorflow func
#
# numpy doesn't work for all ops, make it tf.Tensor()
input_tensors = [
tf.TensorSpec(shape=v.shape,
dtype=tf.as_dtype(v.dtype),
name=utils.node_name(k))
for k, v in feed_dict.items()
]
input_list = [
tf.convert_to_tensor(v,
dtype=tf.as_dtype(v.dtype),
name=utils.node_name(k))
for k, v in feed_dict.items()
]
tf.random.set_seed(1)
result = func(*input_list)
if isinstance(result, (list, tuple)):
# list or tuple
result = [x.numpy() for x in result]
else:
# single result
result = [result.numpy()]
# now make the eager functions a graph
concrete_func = tf.function(func,
input_signature=tuple(input_tensors))
concrete_func = concrete_func.get_concrete_function()
graph_def = from_function(concrete_func,
input_names=list(feed_dict.keys()),
output_names=outputs,
large_model=large_model)
initialized_tables = None
else:
#
# use graph to execute the tensorflow func
#
with tf_session() as sess:
tf_set_random_seed(1)
input_list = []
if not premade_placeholders:
for k, v in clean_feed_dict.items():
input_list.append(
tf_placeholder(name=k,
shape=v.shape,
dtype=tf.as_dtype(v.dtype)))
func(*input_list)
variables_lib.global_variables_initializer().run()
tf_tables_initializer().run()
output_dict = []
for out_name in outputs:
output_dict.append(sess.graph.get_tensor_by_name(out_name))
result = sess.run(output_dict, feed_dict=feed_dict)
graph_def = freeze_session(sess,
input_names=list(feed_dict.keys()),
output_names=outputs)
table_names, key_dtypes, value_dtypes = get_hash_table_info(
graph_def)
initialized_tables = {}
for n, k_dtype, val_dtype in zip(table_names, key_dtypes,
value_dtypes):
h = lookup_ops.hash_table_v2(k_dtype,
val_dtype,
shared_name=n)
k, v = lookup_ops.lookup_table_export_v2(
h, k_dtype, val_dtype)
initialized_tables[n] = (sess.run(k), sess.run(v))
tf_reset_default_graph()
with tf_session() as sess:
tf.import_graph_def(graph_def, name='')
graph_def = tf_optimize(list(feed_dict.keys()),
outputs,
graph_def,
fold_constant=constant_fold)
model_path = os.path.join(
self.test_data_directory,
self._testMethodName + "_after_tf_optimize.pb")
utils.save_protobuf(model_path, graph_def)
self.logger.debug("created file %s", model_path)
return result, graph_def, initialized_tables
def get_node_by_name(self, name):
"""Get node by name."""
ret = self._nodes_by_name.get(name)
if not ret and name:
ret = self._nodes_by_name.get(node_name(name))
return ret
def freeze_and_run_tf(self, func, feed_dict, outputs, as_session,
premade_placeholders, large_model):
np.random.seed(1) # Make it reproducible.
clean_feed_dict = {utils.node_name(k): v for k, v in feed_dict.items()}
if is_tf2() and not as_session:
#
# use eager to execute the tensorflow func
#
# numpy doesn't work for all ops, make it tf.Tensor()
input_tensors = [
tf.TensorSpec(shape=v.shape,
dtype=tf.as_dtype(v.dtype),
name=utils.node_name(k))
for k, v in feed_dict.items()
]
input_list = [
tf.convert_to_tensor(v,
dtype=tf.as_dtype(v.dtype),
name=utils.node_name(k))
for k, v in feed_dict.items()
]
tf.random.set_seed(1)
result = func(*input_list)
if isinstance(result, (list, tuple)):
# list or tuple
result = [x.numpy() for x in result]
else:
# single result
result = [result.numpy()]
# now make the eager functions a graph
concrete_func = tf.function(func,
input_signature=tuple(input_tensors))
concrete_func = concrete_func.get_concrete_function()
graph_def = from_function(concrete_func,
input_names=list(feed_dict.keys()),
output_names=outputs,
large_model=large_model)
initialized_tables = None
else:
#
# use graph to execute the tensorflow func
#
with tf_session() as sess:
tf_set_random_seed(1)
input_list = []
if not premade_placeholders:
for k, v in clean_feed_dict.items():
input_list.append(
tf_placeholder(name=k,
shape=v.shape,
dtype=tf.as_dtype(v.dtype)))
func(*input_list)
variables_lib.global_variables_initializer().run()
tf_tables_initializer().run()
output_dict = []
for out_name in outputs:
output_dict.append(sess.graph.get_tensor_by_name(out_name))
result = sess.run(output_dict, feed_dict=feed_dict)
graph_def = freeze_session(sess,
input_names=list(feed_dict.keys()),
output_names=outputs)
table_info = get_hash_table_info(graph_def)
initialized_tables = {}
for info in table_info:
if info.shared_name is None:
continue
h = lookup_ops.hash_table_v2(info.key_dtype,
info.val_dtype,
shared_name=info.shared_name)
k, v = lookup_ops.lookup_table_export_v2(
h, info.key_dtype, info.val_dtype)
initialized_tables[info.shared_name] = (sess.run(k),
sess.run(v))
tf_reset_default_graph()
with tf_session() as sess:
tf.import_graph_def(graph_def, name='')
graph_def = tf_optimize(list(feed_dict.keys()), outputs,
graph_def)
return result, graph_def, initialized_tables
