def from_graph_def(graph_def, name=None, input_names=None, output_names=None, opset=None, custom_ops=None,
custom_op_handlers=None, custom_rewriter=None, inputs_as_nchw=None, extra_opset=None,
shape_override=None, target=None, large_model=False, tensors_to_rename=None, output_path=None):
"""Returns a ONNX model_proto for a tensorflow graphdef.
Args:
graph_def: the graphdef we want to convert
input_names: list of input names
output_names: list of output names
name: A name for the graph
opset: the opset to be used for the ONNX model, default is the latest
custom_ops: if a model contains ops not recognized by onnx runtime,
you can tag these ops with a custom op domain so that the
runtime can still open the model. Type is a dictionary `{op name: domain}`.
target: list of workarounds applied to help certain platforms
custom_op_handlers: dictionary of custom ops handlers
custom_rewriter: list of custom graph rewriters
extra_opset: list of extra opset's, for example the opset's used by custom ops
shape_override: dict with inputs that override the shapes given by tensorflow
inputs_as_nchw: transpose inputs in list from nchw to nhwc
large_model: use the ONNX external tensor storage format
output_path: save model to output_path
Returns:
An ONNX model_proto and an external_tensor_storage dict.
"""
if not input_names:
raise ValueError("input_names needs to be provided")
if not output_names:
raise ValueError("output_names needs to be provided")
if not name:
name = "unknown"
initialized_tables = None
with tf.device("/cpu:0"):
with tf.Graph().as_default() as tf_graph:
with tf_loader.tf_session(graph=tf_graph) as sess:
tf.import_graph_def(graph_def, name='')
frozen_graph = tf_loader.freeze_session(sess, input_names=input_names, output_names=output_names)
input_names = tf_loader.inputs_without_resource(sess, input_names)
frozen_graph = tf_loader.tf_optimize(input_names, output_names, graph_def)
model_proto, external_tensor_storage = _convert_common(
frozen_graph,
name=name,
continue_on_error=True,
target=target,
opset=opset,
custom_op_handlers=custom_ops,
extra_opset=extra_opset,
shape_override=shape_override,
input_names=input_names,
output_names=output_names,
inputs_as_nchw=inputs_as_nchw,
large_model=large_model,
tensors_to_rename=tensors_to_rename,
initialized_tables=initialized_tables,
output_path=output_path)
return model_proto, external_tensor_storage
def _make_onnx_model():
with tf.compat.v1.Session(graph=tf.Graph()) as session:
input_x = tf.compat.v1.placeholder(dtype=tf.float32,
shape=[None, 2, 3, 4],
name='x')
input_y = tf.compat.v1.placeholder(dtype=tf.float32,
shape=[None, 2, 3, 4],
name='y')
weight = tf.Variable(initial_value=4.2, dtype=tf.float32)
tf.multiply(input_x + input_y, weight, name='z')
session.run(weight.initializer)
frozen_graph_def = tf_loader.freeze_session(session,
input_names=['x:0', 'y:0'],
output_names=['z:0'])
with tf.compat.v1.Session(graph=tf.Graph()) as session:
tf.import_graph_def(frozen_graph_def, name='')
onnx_model = tfonnx.process_tf_graph(
tf_graph=session.graph,
input_names=['x:0', 'y:0'],
output_names=['z:0']).make_model(graph_doc='Test onnx model')
return OnnxModel(model_proto=onnx_model,
input_data_formats=[
DataFormat.CHANNELS_LAST, DataFormat.CHANNELS_FIRST
])
def export_model(self,
export_dir: str,
onnx: Optional[int] = None) -> None:
"""Export tensorflow graph to export_dir for serving."""
if onnx:
try:
import tf2onnx
except ImportError as e:
raise RuntimeError(
"Converting a TensorFlow model to ONNX requires "
"`tf2onnx` to be installed. Install with "
"`pip install tf2onnx`.") from e
with self.get_session().graph.as_default():
signature_def_map = self._build_signature_def()
sd = signature_def_map[
tf1.saved_model.signature_constants.
DEFAULT_SERVING_SIGNATURE_DEF_KEY # noqa: E501
]
inputs = [v.name for k, v in sd.inputs.items()]
outputs = [v.name for k, v in sd.outputs.items()]
from tf2onnx import tf_loader
frozen_graph_def = tf_loader.freeze_session(
self._sess, input_names=inputs, output_names=outputs)
with tf1.Session(graph=tf.Graph()) as session:
tf.import_graph_def(frozen_graph_def, name="")
g = tf2onnx.tfonnx.process_tf_graph(
session.graph,
input_names=inputs,
output_names=outputs,
inputs_as_nchw=inputs,
)
model_proto = g.make_model("onnx_model")
tf2onnx.utils.save_onnx_model(export_dir,
"saved_model",
feed_dict={},
model_proto=model_proto)
else:
with self.get_session().graph.as_default():
signature_def_map = self._build_signature_def()
builder = tf1.saved_model.builder.SavedModelBuilder(export_dir)
builder.add_meta_graph_and_variables(
self.get_session(),
[tf1.saved_model.tag_constants.SERVING],
signature_def_map=signature_def_map,
saver=tf1.summary.FileWriter(export_dir).add_graph(
graph=self.get_session().graph),
)
builder.save()
def load(self, model_path: Union[str, Path], **_) -> Model:
if isinstance(model_path, Path):
model_path = model_path.as_posix()
get_model = load_from_file(model_path, "model", GET_MODEL_FN_NAME)
get_serving_input_receiver_fn = load_from_file(model_path, "model", GET_SERVING_INPUT_RECEIVER_FN)
if get_model is None:
raise RuntimeError(f"Could not find {GET_MODEL_FN_NAME} in {model_path}")
if get_serving_input_receiver_fn is None:
raise RuntimeError(f"Could not find {GET_SERVING_INPUT_RECEIVER_FN} in {model_path}")
model_args = filter_fn_args(self._model_args, fn=get_model)
serving_input_receiver_args = filter_fn_args(self._model_args, fn=get_serving_input_receiver_fn)
session_config = create_session_config(allow_growth=True)
tf.compat.v1.reset_default_graph()
with tf.compat.v1.Session(config=session_config) as sess:
estimator = get_model(**model_args)
serving_input_receiver_fn = get_serving_input_receiver_fn(**serving_input_receiver_args)
input_receiver = serving_input_receiver_fn()
estimator_spec = estimator.model_fn(
features=input_receiver.features,
labels=None,
mode=tf.estimator.ModeKeys.PREDICT,
config=estimator.config,
)
input_tensors_dict = input_receiver.receiver_tensors
output_tensors_dict = estimator_spec.predictions
inputs_dict = {k: tensor2tensor_spec(tensor) for k, tensor in input_tensors_dict.items()}
outputs_dict = {k: tensor2tensor_spec(tensor) for k, tensor in output_tensors_dict.items()}
input_tensor_names = [t.name for t in inputs_dict.values()]
output_tensor_names = [t.name for t in outputs_dict.values()]
graph_saver = estimator_spec.scaffold.saver or tf.compat.v1.train.Saver(sharded=True)
graph_saver.restore(sess, estimator.latest_checkpoint())
input_tensor_names = inputs_without_resource(sess, input_tensor_names)
frozen_graph = freeze_session(sess, input_names=input_tensor_names, output_names=output_tensor_names)
input_tensor_names = remove_redundant_inputs(frozen_graph, input_tensor_names)
tf.compat.v1.reset_default_graph()
with tf.compat.v1.Session(config=estimator.config.session_config):
frozen_graph = tf_optimize(input_tensor_names, output_tensor_names, frozen_graph)
tf.compat.v1.reset_default_graph()
precision = _infer_model_precision(frozen_graph, inputs_dict, outputs_dict)
return Model(frozen_graph, precision, inputs_dict, outputs_dict)
def _from_saved_model_v1(sess, model_path, tag, signatures):
"""
Load tensorflow graph from saved_model.
NOTICE: Modified version from tf2onnx project
"""
wrn_no_tag = "'--tag' not specified for saved_model. Using --tag serve"
wrn_empty_tag = "'--tag' value is empty string. Using tag =[[]]"
if tag is None:
tag = [tf.saved_model.SERVING]
LOGGER.warning(wrn_no_tag)
if tag == "":
tag = [[]]
LOGGER.warning(wrn_empty_tag)
if not isinstance(tag, list):
tag = [tag]
imported = tf.compat.v1.saved_model.loader.load(sess, tag, model_path)
for k in imported.signature_def.keys():
if k.startswith("_"):
# consider signatures starting with '_' private
continue
signatures.append(k)
try:
from tensorflow.contrib.saved_model.python.saved_model import ( # pytype: disable=import-error
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]
inputs = {}
outputs = {}
for k in signatures:
inputs_tensor_info = get_signature_def(imported, k).inputs
for name, input_tensor in inputs_tensor_info.items():
inputs[name] = input_tensor.name
outputs_tensor_info = get_signature_def(imported, k).outputs
for name, output_tensor in outputs_tensor_info.items():
outputs[name] = output_tensor.name
frozen_graph = freeze_session(sess,
input_names=list(inputs.values()),
output_names=list(outputs.values()))
return frozen_graph, inputs, outputs
def _from_keras_tf1(model, input_signature=None, opset=None, custom_ops=None, custom_op_handlers=None,
custom_rewriter=None, inputs_as_nchw=None, extra_opset=None, shape_override=None,
target=None, large_model=False, output_path=None):
"""from_keras for tf 1.15"""
input_names = [t.name for t in model.inputs]
output_names = [t.name for t in model.outputs]
old_out_names = _rename_duplicate_keras_model_names(model)
tensors_to_rename = dict(zip(input_names, model.input_names))
tensors_to_rename.update(zip(output_names, model.output_names))
if old_out_names is not None:
model.output_names = old_out_names
if _is_legacy_keras_model(model):
import keras # pylint: disable=import-outside-toplevel
sess = keras.backend.get_session()
else:
sess = tf.keras.backend.get_session(model.outputs)
with tf.device("/cpu:0"):
frozen_graph, initialized_tables = tf_loader.freeze_session(sess, input_names, output_names, get_tables=True)
with tf.Graph().as_default():
tf.import_graph_def(frozen_graph, name="")
frozen_graph = tf_loader.tf_optimize(input_names, output_names, frozen_graph, False)
model_proto, external_tensor_storage = _convert_common(
frozen_graph,
name=model.name,
continue_on_error=True,
target=target,
opset=opset,
custom_ops=custom_ops,
custom_op_handlers=custom_op_handlers,
custom_rewriter=custom_rewriter,
extra_opset=extra_opset,
shape_override=shape_override,
input_names=input_names,
output_names=output_names,
inputs_as_nchw=inputs_as_nchw,
large_model=large_model,
tensors_to_rename=tensors_to_rename,
initialized_tables=initialized_tables,
output_path=output_path)
return model_proto, external_tensor_storage
def _save_onnx_model(model_file):
with tf.compat.v1.Session(graph=tf.Graph()) as session:
input_x = tf.compat.v1.placeholder(dtype=tf.float32,
shape=[None, 2, 3, 4],
name='x')
input_y = tf.compat.v1.placeholder(dtype=tf.float32,
shape=[None, 2, 3, 4],
name='y')
weight = tf.Variable(initial_value=4.2, dtype=tf.float32)
tf.multiply(input_x + input_y, weight, name='z')
session.run(weight.initializer)
frozen_graph_def = tf_loader.freeze_session(session,
input_names=['x:0', 'y:0'],
output_names=['z:0'])
with tf.compat.v1.Session(graph=tf.Graph()) as session, open(
model_file.name, mode='wb') as graph_file:
tf.import_graph_def(frozen_graph_def, name='')
onnx_model = tfonnx.process_tf_graph(
tf_graph=session.graph,
input_names=['x:0', 'y:0'],
output_names=['z:0']).make_model(graph_doc='Test onnx model')
graph_file.write(onnx_model.SerializeToString())
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 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
