def _run_test_case(self, input_names_with_port, output_names_with_port):
try:
tf.compat.v1.disable_eager_execution()
except: # pylint: disable=bare-except
pass
graph_def = None
with tf_session() as sess:
# freeze graph
origin_graph = sess.graph
variables_lib.global_variables_initializer().run()
output_name_without_port = [
n.split(':')[0] for n in output_names_with_port
]
graph_def = tf.graph_util.convert_variables_to_constants(
sess, sess.graph_def, output_name_without_port)
tf_reset_default_graph()
tf.import_graph_def(graph_def, name='')
# optimize graph
input_tensors = {
i: sess.graph.get_tensor_by_name(i)
for i in input_names_with_port
}
output_tensors = {
i: sess.graph.get_tensor_by_name(i)
for i in output_names_with_port
}
graph_def = tf_optimize(input_tensors, output_tensors, sess.graph_def,
True)
with tf_session() as sess:
if self.config.is_debug_mode:
if not os.path.exists(self.test_data_directory):
os.makedirs(self.test_data_directory)
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)
tf_reset_default_graph()
tf.import_graph_def(graph_def, name='')
with tf_session() as sess:
inferred_graph = infer_shape_for_graph(sess.graph)
# compare each operation
for op in origin_graph.get_operations():
inferred_op = None
try:
inferred_op = inferred_graph.get_operation_by_name(op.name)
except KeyError:
continue
self._compare_shape_for_op(op, inferred_op)
def convert_to_tflite(self, graph_def, feed_dict, outputs):
if not feed_dict:
return None # Can't make TFlite model with no inputs
tf_reset_default_graph()
with tf_session() as sess:
tf.import_graph_def(graph_def, name='')
sess_inputs = [
sess.graph.get_tensor_by_name(k) for k in feed_dict.keys()
]
sess_outputs = [sess.graph.get_tensor_by_name(n) for n in outputs]
converter = tf_lite.TFLiteConverter.from_session(
sess, sess_inputs, sess_outputs)
#converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops.
tf.lite.OpsSet.SELECT_TF_OPS, # enable TensorFlow flex ops.
]
from tensorflow.lite.python.convert import ConverterError
try:
tflite_model = converter.convert()
tflite_path = os.path.join(self.test_data_directory,
self._testMethodName + ".tflite")
dir_name = os.path.dirname(tflite_path)
if dir_name:
os.makedirs(dir_name, exist_ok=True)
with open(tflite_path, 'wb') as f:
f.write(tflite_model)
return tflite_path
except ConverterError:
return None
def test_tensor_data(self):
tensors = {
"empty_tensor": np.array([], dtype=np.float32),
"multi_dim_empty_tensor": np.array([[], []], dtype=np.float32),
"scalar": np.array(1., dtype=np.float32),
"one_item_array": np.array([1.], dtype=np.float32),
"normal_array": np.array([[1., 2.], [2., 3.]], dtype=np.float32)
}
tf_reset_default_graph()
with tf_session() as sess:
for n, data in tensors.items():
tf.constant(data, dtype=tf.float32, name=n)
for tf_node in sess.graph.get_operations():
name = tf_node.name
self.assertTrue(name in tensors.keys())
self.assertTrue("value" in tf_node.node_def.attr)
# convert to onnx tensor value
tensor_value = tf_utils.tf_to_onnx_tensor(
tf_utils.get_tf_node_attr(tf_node, "value"),
name=utils.port_name(tf_node.name)
)
attr = helper.make_attribute("value", tensor_value)
# same as node.get_tensor_value(is_list=False)
actual = numpy_helper.to_array(helper.get_attribute_value(attr))
expected = tensors[name]
self.assertTrue(np.array_equal(expected, actual))
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 test_parse_tflite_graph(self):
def func(a, b, c):
alpha = tf.constant(1.1, dtype=tf.float32)
beta = tf.constant(2.3, dtype=tf.float32)
mul1 = tf.multiply(alpha, tf.matmul(a, b))
mul2 = tf.multiply(beta, c)
x_ = mul1 + mul2
return tf.identity(x_, name="output")
inp_shapes = [[2, 3], [3, 1], [2, 1]]
inp_dtypes = [tf.float32, tf.float32, tf.float32]
names = ['a', 'b', 'c']
names_with_port = ['a:0', 'b:0', 'c:0']
output_names = ['output']
output_names_with_port = ['output:0']
input_tensors = [tf.TensorSpec(shape=s, dtype=d, name=n) for s, d, n in zip(inp_shapes, inp_dtypes, names)]
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=names_with_port,
output_names=output_names_with_port)
with tf_session() as sess:
tf.import_graph_def(graph_def, name='')
sess_inputs = [sess.graph.get_tensor_by_name(k) for k in names_with_port]
sess_outputs = [sess.graph.get_tensor_by_name(n) for n in output_names_with_port]
converter = tf.compat.v1.lite.TFLiteConverter.from_session(sess, sess_inputs, sess_outputs)
tflite_model = converter.convert()
tflite_path = os.path.join(self.test_data_directory, self._testMethodName + ".tflite")
dir_name = os.path.dirname(tflite_path)
tflite_model = converter.convert()
os.makedirs(dir_name, exist_ok=True)
with open(tflite_path, 'wb') as f:
f.write(tflite_model)
tflite_graphs, opcodes_map, model, tensor_shapes = read_tflite_model(tflite_path)
self.assertEqual(1, len(tflite_graphs))
onnx_nodes, op_cnt, attr_cnt, output_shapes, dtypes, inputs, outputs, _ = \
parse_tflite_graph(tflite_graphs[0], opcodes_map, model, tensor_shapes_override=tensor_shapes)
self.assertEqual(2, op_cnt['MUL'])
self.assertEqual(1, op_cnt['ADD'])
self.assertEqual(1, op_cnt['FULLY_CONNECTED'])
self.assertEqual(1, attr_cnt['WeightsFormat'])
self.assertEqual(names, inputs)
self.assertEqual(output_names, outputs)
for name, shape, dtype in zip(names, inp_shapes, inp_dtypes):
self.assertEqual(shape, output_shapes[name])
self.assertEqual(dtype, dtypes[name])
self.assertTrue(len(onnx_nodes) >= 4)
def get_output_shapes(node_def, input_dtypes, input_shapes, inp_consts):
"""Returns a list of the output shapes of an op. input_dtypes should be tf dtypes."""
from tf2onnx.tf_loader import tf_session, tf_placeholder # pylint: disable=import-outside-toplevel
if node_def.op in ["Prelu", "Enter"]:
return [input_shapes[0]]
if node_def.op == "Merge":
# Find the first non-None shape (if it exists) and return it
non_none = ([t for t in input_shapes if t is not None] + [None])[0]
# The second output of merge is a scalar int indicating which input was selected
return [non_none, []]
if node_def.op == "Placeholder":
shape = None
if 'shape' in node_def.attr:
shape = [d.size for d in node_def.attr['shape'].shape.dim]
shape = [None if d == -1 else d for d in shape]
if len(shape) == 0:
# According to TF docs, "If the shape has 0 dimensions, the shape is unconstrained."
shape = None
return [shape]
del node_def.input[:]
node_def.name = "node"
if "_class" in node_def.attr:
# Remove colocation information (list of nodes tf wants computed on same device)
del node_def.attr["_class"]
g = tf.Graph()
with g.as_default():
for i, (dtype, shape,
const) in enumerate(zip(input_dtypes, input_shapes,
inp_consts)):
inp = "input" + str(i)
if const is None:
if shape is not None and -1 in shape:
shape = [d if d != -1 else None for d in shape]
tf_placeholder(dtype, name=inp, shape=shape)
else:
tf.constant(const, dtype, name=inp)
node_def.input.append(inp)
mini_graph_def = g.as_graph_def()
mini_graph_def.node.append(node_def)
g2 = tf.Graph()
with g2.as_default():
with tf_session() as sess:
tf.import_graph_def(mini_graph_def, name='')
node = sess.graph.get_operation_by_name("node")
outputs_shapes = [
tf_utils.get_tf_tensor_shape(out) for out in node.outputs
]
return outputs_shapes
def convert_to_onnx(self, graph_def, inputs, outputs):
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
with tf_loader.tf_session(graph=tf_graph):
g = process_tf_graph(tf_graph,
continue_on_error=False,
target=",".join(constants.DEFAULT_TARGET),
opset=11,
input_names=inputs,
output_names=outputs,
inputs_as_nchw=None)
onnx_graph = optimizer.optimize_graph(g)
model_proto = onnx_graph.make_model("converted from {}".format(
self._tf_file))
return model_proto
def __init__(self, saved_model_path, legacy_plugins=False):
"""
Constructor of the EfficientDet Graph Surgeon object, to do the conversion of an EfficientDet TF saved model
to an ONNX-TensorRT parsable model.
:param saved_model_path: The path pointing to the TensorFlow saved model to load.
:param legacy_plugins: If using TensorRT version < 8.0.1, set this to True to use older (but slower) plugins.
"""
saved_model_path = os.path.realpath(saved_model_path)
assert os.path.exists(saved_model_path)
# Use tf2onnx to convert saved model to an initial ONNX graph.
graph_def, inputs, outputs = tf_loader.from_saved_model(
saved_model_path, None, None, "serve", ["serving_default"])
log.info("Loaded saved model from {}".format(saved_model_path))
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name="")
with tf_loader.tf_session(graph=tf_graph):
onnx_graph = tfonnx.process_tf_graph(tf_graph,
input_names=inputs,
output_names=outputs,
opset=11)
onnx_model = optimizer.optimize_graph(onnx_graph).make_model(
"Converted from {}".format(saved_model_path))
self.graph = gs.import_onnx(onnx_model)
assert self.graph
log.info("TF2ONNX graph created successfully")
# Fold constants via ONNX-GS that TF2ONNX may have missed
self.graph.fold_constants()
# Try to auto-detect by finding if nodes match a specific name pattern expected for either of the APIs.
self.api = None
if len(
[node
for node in self.graph.nodes if "class_net/" in node.name]) > 0:
self.api = "AutoML"
elif len([
node for node in self.graph.nodes
if "/WeightSharedConvolutionalClassHead/" in node.name
]) > 0:
self.api = "TFOD"
assert self.api
log.info("Graph was detected as {}".format(self.api))
self.batch_size = None
self.legacy_plugins = legacy_plugins
def read_tf_node_def_attrs(node_def, input_dtypes, input_shapes):
"""Given a tf node def, returns a dict of attribute names to values"""
from tf2onnx.tf_loader import tf_session, tf_placeholder # pylint: disable=import-outside-toplevel
del node_def.input[:]
node_def.name = "node"
# read_tf_node_attrs uses some tf methods that require the node to be loaded into a valid TF graph
g = tf.Graph()
with g.as_default():
for i, (dtype, shape) in enumerate(zip(input_dtypes, input_shapes)):
inp = "input" + str(i)
tf_placeholder(dtype, name=inp, shape=shape)
node_def.input.append(inp)
mini_graph_def = g.as_graph_def()
mini_graph_def.node.append(node_def)
g2 = tf.Graph()
with g2.as_default():
with tf_session() as sess:
tf.import_graph_def(mini_graph_def, name='')
node = sess.graph.get_operation_by_name("node")
return read_tf_node_attrs(node)
def run_test(self,
name,
backend="caffe2",
onnx_file=None,
opset=None,
extra_opset=None,
perf=None,
fold_const=None):
"""Run complete test against backend."""
self.perf = perf
# get the model
if self.url:
_, dir_name = self.download_model()
logger.info("Downloaded to %s", dir_name)
model_path = os.path.join(
dir_name, self.local) if self.local != "." else dir_name
else:
model_path = self.local
logger.info("Load model from %s", model_path)
input_names = list(self.input_names.keys())
outputs = self.output_names
if self.model_type in ["checkpoint"]:
graph_def, input_names, outputs = tf_loader.from_checkpoint(
model_path, input_names, outputs)
elif self.model_type in ["saved_model"]:
loaded = tf_loader.from_saved_model(
model_path,
input_names,
outputs,
self.tag,
self.signatures,
self.concrete_function,
self.large_model,
return_concrete_func=self.large_model)
if self.large_model:
# Must maintain ref to imported since concrete_func uses weak refs
# pylint: disable=unused-variable
graph_def, input_names, outputs, concrete_func, imported = loaded
else:
graph_def, input_names, outputs = loaded
elif self.model_type in ["keras"]:
graph_def, input_names, outputs = tf_loader.from_keras(
model_path, input_names, outputs)
else:
graph_def, input_names, outputs = tf_loader.from_graphdef(
model_path, input_names, outputs)
if utils.is_debug_mode():
utils.save_protobuf(
os.path.join(TEMP_DIR, name + "_after_tf_optimize.pb"),
graph_def)
if self.large_model:
inputs = {}
for k in input_names:
v = self.input_names[k]
inputs[k.split(":")[0]] = tf.constant(self.make_input(v))
tf_func = tf.function(concrete_func)
logger.info("Running TF")
tf_results_d = tf_func(**inputs)
if self.structured_outputs is None:
tf_results = list(tf_results_d.values())
else:
tf_results = [
tf_results_d[output] for output in self.structured_outputs
]
if self.perf:
logger.info("Running TF perf")
start = time.time()
for _ in range(PERFITER):
_ = concrete_func(**inputs)
self.tf_runtime = time.time() - start
logger.info("TensorFlow OK")
inputs = {}
shape_override = {}
tf_reset_default_graph()
from tf2onnx.tf_utils import compress_graph_def
const_node_values = None
with tf.Graph().as_default() as tf_graph:
if self.large_model:
const_node_values = compress_graph_def(graph_def)
tf.import_graph_def(graph_def, name='')
with tf_session(graph=tf_graph) as sess:
# create the input data
for k in input_names:
v = self.input_names[k]
t = sess.graph.get_tensor_by_name(k)
expected_dtype = tf.as_dtype(t.dtype).name
if isinstance(v, six.text_type) and v.startswith("np."):
np_value = eval(v) # pylint: disable=eval-used
if expected_dtype != np_value.dtype:
logger.warning(
"dtype mismatch for input %s: expected=%s, actual=%s",
k, expected_dtype, np_value.dtype)
inputs[k] = np_value.astype(expected_dtype)
else:
inputs[k] = self.make_input(v).astype(expected_dtype)
if self.force_input_shape:
for k, v in inputs.items():
shape_override[k] = list(v.shape)
# run the model with tensorflow
if self.skip_tensorflow:
logger.info("TensorFlow SKIPPED")
elif not self.large_model:
tf_results = self.run_tensorflow(sess, inputs)
logger.info("TensorFlow OK")
model_proto = None
if self.skip_conversion:
if self.large_model:
external_tensor_storage = ExternalTensorStorage()
model_proto = utils.model_proto_from_zip(
self.converted_model, external_tensor_storage)
else:
external_tensor_storage = None
model_proto = utils.model_proto_from_file(self.converted_model)
logger.info("ONNX loaded from file")
else:
try:
# convert model to onnx
onnx_graph = self.to_onnx(sess.graph,
opset=opset,
extra_opset=extra_opset,
shape_override=shape_override,
input_names=inputs.keys(),
const_node_values=const_node_values)
onnx_graph = optimizer.optimize_graph(onnx_graph)
print("ONNX", onnx_graph.dump_node_statistics())
external_tensor_storage = ExternalTensorStorage(
) if self.large_model else None
model_proto = onnx_graph.make_model(
"converted from tf2onnx",
external_tensor_storage=external_tensor_storage)
logger.info("To_ONNX, OK")
if onnx_file:
self.create_onnx_file(name, model_proto, inputs, onnx_file,
external_tensor_storage)
if self.converted_model:
if self.large_model:
utils.save_onnx_zip(self.converted_model, model_proto,
external_tensor_storage)
else:
utils.save_protobuf(self.converted_model, model_proto)
logger.info("Created %s", self.converted_model)
except Exception:
logger.error("To_ONNX FAIL", exc_info=1)
return False
try:
onnx_results = None
if backend == "caffe2":
onnx_results = self.run_caffe2(name, model_proto, inputs)
elif backend == "onnxruntime":
onnx_results = self.run_onnxruntime(name, model_proto, inputs,
external_tensor_storage)
else:
raise ValueError("unknown backend")
logger.info("Run_ONNX OK")
try:
if self.skip_tensorflow:
logger.info("Results: skipped tensorflow")
else:
if self.check_only_shape:
for tf_res, onnx_res in zip(tf_results, onnx_results):
np.testing.assert_array_equal(
tf_res.shape, onnx_res.shape)
else:
for tf_res, onnx_res in zip(tf_results, onnx_results):
np.testing.assert_allclose(tf_res,
onnx_res,
rtol=self.rtol,
atol=self.atol)
logger.info("Results: OK")
return True
except Exception:
logger.error("Results", exc_info=1)
except Exception:
logger.error("Run_ONNX FAIL", exc_info=1)
return False
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,
premade_placeholders=False):
test_tf = not self.config.skip_tf_tests
test_tflite = not self.config.skip_tflite_tests
run_tfl_consistency_test = test_tf and test_tflite and self.config.run_tfl_consistency_test
# 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()
if tf_lite is None:
test_tflite = False
g = None
expected, graph_def, initialized_tables = \
self.freeze_and_run_tf(func, feed_dict, output_names_with_port, as_session,
premade_placeholders, large_model, constant_fold)
if test_tflite:
tflite_path = self.convert_to_tflite(graph_def, feed_dict,
output_names_with_port)
test_tflite = tflite_path is not None
if test_tf:
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='')
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,
initialized_tables=initialized_tables,
**process_args)
g = optimizer.optimize_graph(g, catch_errors=False)
actual = self.run_backend(g, output_names_with_port,
onnx_feed_dict, large_model)
self.assert_results_equal(expected, actual, rtol, atol,
check_value, check_shape, check_dtype)
self.assert_shapes_correct(g, self.config.allow_missing_shapes,
not self.config.skip_onnx_checker)
if graph_validator:
self.assertTrue(graph_validator(g))
if test_tflite:
tfl_results, tfl_outputs = self.run_tflite(tflite_path, feed_dict)
test_tflite = tfl_results is not None
if test_tflite:
if run_tfl_consistency_test:
self.assert_results_equal(expected, tfl_results, rtol, atol,
check_value, check_shape,
check_dtype)
tfl_process_args = process_args.copy()
if 'inputs_as_nchw' in tfl_process_args:
nchw_inps_with_port = tfl_process_args['inputs_as_nchw']
tfl_process_args['inputs_as_nchw'] = [
i.split(':')[0] for i in nchw_inps_with_port
]
input_names_without_port = [
inp.split(':')[0] for inp in feed_dict.keys()
]
g = process_tf_graph(None,
opset=self.config.opset,
input_names=input_names_without_port,
output_names=tfl_outputs,
target=self.config.target,
tflite_path=tflite_path,
**tfl_process_args)
g = optimizer.optimize_graph(g)
onnx_feed_dict_without_port = {
k.split(':')[0]: v
for k, v in onnx_feed_dict.items()
}
onnx_from_tfl_res = self.run_backend(g,
tfl_outputs,
onnx_feed_dict_without_port,
postfix="_from_tflite")
self.assert_results_equal(tfl_results, onnx_from_tfl_res, rtol,
atol, check_value, check_shape,
check_dtype)
self.assert_shapes_correct(g, self.config.allow_missing_shapes,
not self.config.skip_onnx_checker)
if graph_validator:
self.assertTrue(graph_validator(g))
if g is None:
raise unittest.SkipTest("Both tf and tflite marked to skip")
return g
def tf2onnx_flow(pb_path: str, test_mode =False) -> onnx.ModelProto:
"""Convert frozen graph pb file into onnx
Args:
pb_path (str): input pb file path
test_mode (bool, optional): test mode. Defaults to False.
Raises:
Exception: invalid input file
Returns:
onnx.ModelProto: converted onnx
"""
TF2ONNX_VERSION = int(tf2onnx.version.version.replace('.', ''))
if 160 <= TF2ONNX_VERSION:
from tf2onnx import tf_loader
else:
from tf2onnx import loader as tf_loader
if pb_path[-3:] == '.pb':
model_name = pb_path.split('/')[-1][:-3]
# always reset tensorflow session at begin
tf.reset_default_graph()
with tf.Session() as sess:
with gfile.FastGFile(pb_path, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
sess.graph.as_default()
tf.import_graph_def(graph_def, name='')
if 160 <= int(tf2onnx.version.version.replace('.', '')):
onnx_nodes, op_cnt, attr_cnt, output_shapes, dtypes, functions = tf2onnx.tf_utils.tflist_to_onnx(
sess.graph,
{})
else:
onnx_nodes, op_cnt, attr_cnt, output_shapes, dtypes = tf2onnx.tfonnx.tflist_to_onnx(
sess.graph.get_operations(),
{})
for n in onnx_nodes:
if len(n.output) == 0:
onnx_nodes.remove(n)
# find inputs and outputs of graph
nodes_inputs = set()
nodes_outputs = set()
for n in onnx_nodes:
if n.op_type == 'Placeholder':
continue
for input in n.input:
nodes_inputs.add(input)
for output in n.output:
nodes_outputs.add(output)
graph_input_names = set()
for input_name in nodes_inputs:
if input_name not in nodes_outputs:
graph_input_names.add(input_name)
graph_output_names = set()
for n in onnx_nodes:
if n.input and n.input[0] not in nodes_outputs:
continue
if len(n.output) == 0:
n.output.append(n.name + ':0')
graph_output_names.add(n.output[0])
else:
output_name = n.output[0]
if (output_name not in nodes_inputs) and (0 < len(n.input)):
graph_output_names.add(output_name)
logging.info('Model Inputs: %s', str(list(graph_input_names)))
logging.info('Model Outputs: %s', str(list(graph_output_names)))
graph_def, inputs, outputs = tf_loader.from_graphdef(model_path=pb_path,
input_names=list(graph_input_names),
output_names=list(graph_output_names))
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
if 160 <= TF2ONNX_VERSION:
with tf_loader.tf_session(graph=tf_graph):
onnx_graph = tf2onnx.tfonnx.process_tf_graph(tf_graph=tf_graph,
input_names=inputs,
output_names=outputs,
opset=11)
else:
with tf.Session(graph=tf_graph):
onnx_graph = tf2onnx.tfonnx.process_tf_graph(tf_graph=tf_graph,
input_names=inputs,
output_names=outputs,
opset=11)
# Optimize with tf2onnx.optimizer
onnx_graph = tf2onnx.optimizer.optimize_graph(onnx_graph)
model_proto = onnx_graph.make_model(model_name)
# Make tf2onnx output compatible with the spec. of onnx.utils.polish_model
replacing.replace_initializer_with_Constant(model_proto.graph)
model_proto = onnx.utils.polish_model(model_proto)
else:
raise Exception('expect .pb file as input, but got "' + str(pb_path) + '"')
# rename
m = model_proto
m = combo.preprocess(m)
m = combo.common_optimization(m)
m = combo.tensorflow_optimization(m)
m = combo.postprocess(m)
if not test_mode:
g = m.graph
eliminating.eliminate_shape_changing_after_input(g)
m = onnx.utils.polish_model(m)
return m
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 compute_const_folding_using_tf(g, const_node_values, graph_outputs):
"""Find nodes with constant inputs and compute their values using TF"""
if const_node_values is None:
const_node_values = {}
graph_outputs = set(graph_outputs)
from tf2onnx.tf_loader import tf_session, tf_placeholder # pylint: disable=import-outside-toplevel
ops = g.get_operations()
outputs_to_values = {}
outputs_to_dtypes = {}
outputs_to_shapes = {}
shape_node_outputs = {}
def is_small_shape(x):
return np.product(x) <= 1000
def is_huge_shape(x):
return np.product(x) >= 1000000
for node in ops:
# Load values of constants. Use const_node_values if possible
if node.type in ["Const", "ConstV2"]:
tensor = node.node_def.attr["value"].tensor
if node.name in const_node_values:
tensor.tensor_content = const_node_values[node.name]
outputs_to_values[node.outputs[0].name] = get_tf_tensor_data(
tensor)
outputs_to_dtypes[node.outputs[0].name] = node.outputs[0].dtype
for out in node.outputs:
outputs_to_shapes[out.name] = get_tf_tensor_shape(out)
for node in ops:
if node.type == "Shape":
shape = outputs_to_shapes.get(node.inputs[0].name)
if shape is not None:
shape_node_outputs[node.outputs[0].name] = shape
unneeded_outputs = set()
progress = True
while progress:
progress = False
for node in ops:
# Find ops with constant inputs and compute their values
input_names = [i.name for i in node.inputs]
output_names = [i.name for i in node.outputs]
if node.type == 'StridedSlice' and input_names[0] in shape_node_outputs \
and output_names[0] not in outputs_to_values:
shape = shape_node_outputs[input_names[0]]
i = get_index_from_strided_slice_of_shape(
node, outputs_to_values)
if i is not None and 0 <= i < len(
shape) and shape[i] is not None:
np_dtype = map_onnx_to_numpy_type(
map_tf_dtype(node.outputs[0].dtype))
outputs_to_values[output_names[0]] = np.array(
shape[i], dtype=np_dtype)
outputs_to_dtypes[
node.outputs[0].name] = node.outputs[0].dtype
progress = True
can_fold = node.type not in [
'Enter', 'Placeholder', 'PlaceholderWithDefault'
]
can_fold = can_fold and len(input_names) > 0 and all(
inp in outputs_to_values for inp in input_names)
# We can only fold nodes with a single output
can_fold = can_fold and len(
output_names) == 1 and output_names[0] not in outputs_to_values
# Skip if value already computed, used, and discarded
can_fold = can_fold and output_names[
0] not in unneeded_outputs and output_names[
0] not in graph_outputs
if can_fold:
# Make a mini graph containing just the node to fold
g2 = tf.Graph()
with g2.as_default():
for inp in input_names:
tf_placeholder(outputs_to_dtypes[inp],
name=inp.split(':')[0])
mini_graph_def = g2.as_graph_def()
mini_graph_def.node.append(node.node_def)
g3 = tf.Graph()
with g3.as_default():
feed_dict = {}
inp_shapes = []
for inp in input_names:
inp_np = outputs_to_values[inp]
feed_dict[inp] = inp_np
inp_shapes.append(inp_np.shape)
try:
with tf_session() as sess:
tf.import_graph_def(mini_graph_def, name='')
results = sess.run(output_names,
feed_dict=feed_dict)
if is_huge_shape(results[0].shape) and all(
is_small_shape(inp) for inp in inp_shapes):
logger.debug(
"Skipping folding of node %s since result shape %s is much larger "
"than input shapes %s", node.name,
results[0].shape, inp_shapes)
else:
outputs_to_values[output_names[0]] = results[0]
outputs_to_dtypes[
output_names[0]] = node.outputs[0].dtype
progress = True
except Exception: # pylint: disable=broad-except
logger.debug("Could not fold node %s", node.name)
unneeded_outputs.update(outputs_to_values.keys())
for node in ops:
# Mark values we need to keep
input_names = [i.name for i in node.inputs]
output_names = [i.name for i in node.outputs]
if len(output_names) == 1 and output_names[0] in outputs_to_values:
continue
for i in input_names:
if i in unneeded_outputs:
unneeded_outputs.remove(i)
for node in unneeded_outputs:
# Remove unneeded values to prevent memory usage explosion
if node in outputs_to_values:
del outputs_to_values[node]
del outputs_to_dtypes[node]
for node in ops:
# We don't need the constants any more
if node.type in ["Const", "ConstV2"
] and node.outputs[0].name in outputs_to_values:
del outputs_to_values[node.outputs[0].name]
del outputs_to_dtypes[node.outputs[0].name]
logger.info("Computed %d values for constant folding",
len(outputs_to_values))
return outputs_to_values, outputs_to_dtypes
def main(args):
# Load saved model
saved_model_path = os.path.realpath(args.saved_model)
assert os.path.isdir(saved_model_path)
graph_def, inputs, outputs = tf_loader.from_saved_model(
saved_model_path, None, None, "serve", ["serving_default"])
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name="")
with tf_loader.tf_session(graph=tf_graph):
onnx_graph = tfonnx.process_tf_graph(tf_graph,
input_names=inputs,
output_names=outputs,
opset=11)
onnx_model = optimizer.optimize_graph(onnx_graph).make_model(
"Converted from {}".format(saved_model_path))
graph = gs.import_onnx(onnx_model)
assert graph
print()
print("ONNX graph created successfully")
# Set the I/O tensor shapes
graph.inputs[0].shape[0] = args.batch_size
graph.outputs[0].shape[0] = args.batch_size
if args.input_size and args.input_size > 0:
if graph.inputs[0].shape[3] == 3:
# Format NHWC
graph.inputs[0].shape[1] = args.input_size
graph.inputs[0].shape[2] = args.input_size
elif graph.inputs[0].shape[1] == 3:
# Format NCHW
graph.inputs[0].shape[2] = args.input_size
graph.inputs[0].shape[3] = args.input_size
print("ONNX input named '{}' with shape {}".format(graph.inputs[0].name,
graph.inputs[0].shape))
print("ONNX output named '{}' with shape {}".format(
graph.outputs[0].name, graph.outputs[0].shape))
for i in range(4):
if type(graph.inputs[0].shape[i]
) != int or graph.inputs[0].shape[i] <= 0:
print(
"The input shape of the graph is invalid, try overriding it by giving a fixed size with --input_size"
)
sys.exit(1)
# Fix Clip Nodes (ReLU6)
for node in [n for n in graph.nodes if n.op == "Clip"]:
for input in node.inputs[1:]:
# In TensorRT, the min/max inputs on a Clip op *must* have fp32 datatype
input.values = np.float32(input.values)
# Run tensor shape inference
graph.cleanup().toposort()
model = shape_inference.infer_shapes(gs.export_onnx(graph))
graph = gs.import_onnx(model)
# Save updated model
graph.cleanup().toposort()
model = gs.export_onnx(graph)
onnx_path = os.path.realpath(args.onnx)
os.makedirs(os.path.dirname(onnx_path), exist_ok=True)
onnx.save(model, onnx_path)
engine_path = os.path.join(os.path.dirname(onnx_path), "engine.trt")
print("ONNX model saved to {}".format(onnx_path))
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
def run_test(self,
name,
backend="caffe2",
onnx_file=None,
opset=None,
extra_opset=None,
perf=None,
fold_const=None):
"""Run complete test against backend."""
self.perf = perf
# get the model
if self.url:
_, dir_name = self.download_model()
logger.info("Downloaded to %s", dir_name)
model_path = os.path.join(dir_name, self.local)
else:
model_path = self.local
logger.info("Load model from %s", model_path)
input_names = list(self.input_names.keys())
outputs = self.output_names
if self.model_type in ["checkpoint"]:
graph_def, input_names, outputs = tf_loader.from_checkpoint(
model_path, input_names, outputs)
elif self.model_type in ["saved_model"]:
graph_def, input_names, outputs = tf_loader.from_saved_model(
model_path, input_names, outputs)
elif self.model_type in ["keras"]:
graph_def, input_names, outputs = tf_loader.from_keras(
model_path, input_names, outputs)
else:
graph_def, input_names, outputs = tf_loader.from_graphdef(
model_path, input_names, outputs)
if utils.is_debug_mode():
utils.save_protobuf(
os.path.join(TEMP_DIR, name + "_after_tf_optimize.pb"),
graph_def)
inputs = {}
shape_override = {}
tf_reset_default_graph()
g = tf.import_graph_def(graph_def, name='')
# with tf_session(config=tf.ConfigProto(allow_soft_placement=True), graph=g) as sess:
with tf_session(graph=g) as sess:
# create the input data
for k in input_names:
v = self.input_names[k]
t = sess.graph.get_tensor_by_name(k)
expected_dtype = tf.as_dtype(t.dtype).name
if isinstance(v, six.text_type) and v.startswith("np."):
np_value = eval(v) # pylint: disable=eval-used
if expected_dtype != np_value.dtype:
logger.warning(
"dtype mismatch for input %s: expected=%s, actual=%s",
k, expected_dtype, np_value.dtype)
inputs[k] = np_value.astype(expected_dtype)
else:
inputs[k] = self.make_input(v).astype(expected_dtype)
if self.force_input_shape:
for k, v in inputs.items():
shape_override[k] = list(v.shape)
# run the model with tensorflow
if self.skip_tensorflow:
logger.info("TensorFlow SKIPPED")
else:
tf_results = self.run_tensorflow(sess, inputs)
logger.info("TensorFlow OK")
model_proto = None
try:
# convert model to onnx
onnx_graph = self.to_onnx(sess.graph,
opset=opset,
extra_opset=extra_opset,
shape_override=shape_override,
input_names=inputs.keys())
onnx_graph = optimizer.optimize_graph(onnx_graph)
model_proto = onnx_graph.make_model("converted from tf2onnx")
logger.info("To_ONNX, OK")
if onnx_file:
self.create_onnx_file(name, model_proto, inputs, onnx_file)
except Exception:
logger.error("To_ONNX FAIL", exc_info=1)
return False
try:
onnx_results = None
if backend == "caffe2":
onnx_results = self.run_caffe2(name, model_proto, inputs)
elif backend == "onnxruntime":
onnx_results = self.run_onnxruntime(name, model_proto, inputs)
else:
raise ValueError("unknown backend")
logger.info("Run_ONNX OK")
try:
if self.skip_tensorflow:
logger.info("Results: skipped tensorflow")
else:
if self.check_only_shape:
for tf_res, onnx_res in zip(tf_results, onnx_results):
np.testing.assert_array_equal(
tf_res.shape, onnx_res.shape)
else:
for tf_res, onnx_res in zip(tf_results, onnx_results):
np.testing.assert_allclose(tf_res,
onnx_res,
rtol=self.rtol,
atol=self.atol)
logger.info("Results: OK")
return True
except Exception:
logger.error("Results", exc_info=1)
except Exception:
logger.error("Run_ONNX FAIL", exc_info=1)
return False
def main():
args = get_args()
logging.basicConfig(level=logging.get_verbosity_level(args.verbose))
if args.debug:
utils.set_debug_mode(True)
logger = logging.getLogger(constants.TF2ONNX_PACKAGE_NAME)
extra_opset = args.extra_opset or []
custom_ops = {}
initialized_tables = None
if args.custom_ops:
using_tf_opset = False
for op in args.custom_ops.split(","):
if ":" in op:
op, domain = op.split(":")
else:
# default custom ops for tensorflow-onnx are in the "tf" namespace
using_tf_opset = True
domain = constants.TENSORFLOW_OPSET.domain
custom_ops[op] = (make_default_custom_op_handler(domain), [])
if using_tf_opset:
extra_opset.append(constants.TENSORFLOW_OPSET)
# get the frozen tensorflow model from graphdef, checkpoint or saved_model.
if args.graphdef:
graph_def, inputs, outputs = tf_loader.from_graphdef(
args.graphdef, args.inputs, args.outputs)
model_path = args.graphdef
if args.checkpoint:
graph_def, inputs, outputs = tf_loader.from_checkpoint(
args.checkpoint, args.inputs, args.outputs)
model_path = args.checkpoint
if args.saved_model:
graph_def, inputs, outputs, initialized_tables = tf_loader.from_saved_model(
args.saved_model,
args.inputs,
args.outputs,
args.tag,
args.signature_def,
args.concrete_function,
args.large_model,
return_initialized_tables=True)
model_path = args.saved_model
if args.keras:
graph_def, inputs, outputs = tf_loader.from_keras(
args.keras, args.inputs, args.outputs)
model_path = args.keras
if args.verbose:
logger.info("inputs: %s", inputs)
logger.info("outputs: %s", outputs)
with tf.Graph().as_default() as tf_graph:
const_node_values = None
if args.large_model:
const_node_values = compress_graph_def(graph_def)
if args.output_frozen_graph:
utils.save_protobuf(args.output_frozen_graph, graph_def)
tf.import_graph_def(graph_def, name='')
with tf_loader.tf_session(graph=tf_graph):
g = process_tf_graph(tf_graph,
continue_on_error=args.continue_on_error,
target=args.target,
opset=args.opset,
custom_op_handlers=custom_ops,
extra_opset=extra_opset,
shape_override=args.shape_override,
input_names=inputs,
output_names=outputs,
inputs_as_nchw=args.inputs_as_nchw,
ignore_default=args.ignore_default,
use_default=args.use_default,
const_node_values=const_node_values,
initialized_tables=initialized_tables)
onnx_graph = optimizer.optimize_graph(g)
tensor_storage = ExternalTensorStorage() if args.large_model else None
model_proto = onnx_graph.make_model("converted from {}".format(model_path),
external_tensor_storage=tensor_storage)
# write onnx graph
logger.info("")
logger.info("Successfully converted TensorFlow model %s to ONNX",
model_path)
if args.output:
if args.large_model:
utils.save_onnx_zip(args.output, model_proto, tensor_storage)
logger.info(
"Zipped ONNX model is saved at %s. Unzip before opening in onnxruntime.",
args.output)
else:
utils.save_protobuf(args.output, model_proto)
logger.info("ONNX model is saved at %s", args.output)
else:
logger.info(
"To export ONNX model to file, please run with `--output` option")
def _run_pb_gen():
'''Load Model from model.py file'''
inference_model = model.Model(is_training=False,
seq_length=FLAGS.seq_length,
batch_size=FLAGS.batch_size,
img_height=FLAGS.img_height,
img_width=FLAGS.img_width)
with tf.compat.v1.Session() as sess:
'''Initialize Variables in Model'''
init_op = tf.compat.v1.global_variables_initializer()
'''Start Session'''
sess.run(init_op)
'''Get Graph Def'''
graph_def = sess.graph.as_graph_def()
'''Extract Inputs'''
inputs = []
for op in sess.graph.get_operations():
if op.type == "Placeholder":
inputs.append(op.name)
'''Extract Outputs'''
name_list = []
exclsv_list = []
for node in graph_def.node:
name_list.append(node.name)
exclsv_list.extend(node.input)
outputs = list(set(name_list) - set(exclsv_list))
outputs = ['depth_prediction/depth_prediction/truediv']
'''Fix Nodes'''
'''See: https://github.com/onnx/tensorflow-onnx/issues/77'''
for node in graph_def.node:
if node.op == 'RefSwitch':
node.op = 'Switch'
for index in range(len(node.input)):
if 'moving_' in node.input[index]:
node.input[index] = node.input[index] + '/read'
elif node.op == 'AssignSub':
node.op = 'Sub'
if 'use_locking' in node.attr: del node.attr['use_locking']
elif node.op == 'AssignAdd':
node.op = 'Add'
if 'use_locking' in node.attr: del node.attr['use_locking']
elif node.op == 'Assign':
node.op = 'Identity'
if 'use_locking' in node.attr: del node.attr['use_locking']
if 'validate_shape' in node.attr:
del node.attr['validate_shape']
if len(node.input) == 2:
# input0: ref: Should be from a Variable node. May be uninitialized.
# input1: value: The value to be assigned to the variable.
node.input[0] = node.input[1]
del node.input[1]
elif node.op == 'L2Loss':
node.op = 'Abs'
'''Sub Graph Extraction'''
needed_names = [tf2onnx.utils.node_name(i) for i in inputs
] + [tf2onnx.utils.node_name(i) for i in outputs]
sub_graph = tf.compat.v1.graph_util.extract_sub_graph(
graph_def, needed_names)
'''Freezing Graph (Necessary before Making ONNX Graph)'''
frozen_graph = freeze_session(sess, sub_graph, output_names=outputs)
frozen_graph = tf.graph_util.remove_training_nodes(frozen_graph)
with open("frozen.pb", "wb") as f:
f.write(frozen_graph.SerializeToString())
'''Graph_Def to Graph Conversion'''
tf_reset_default_graph()
graph = tf.import_graph_def(frozen_graph, name='')
with tf_session(graph=graph) as sess:
'''Extract Inputs'''
inputs = []
for op in sess.graph.get_operations():
if op.type == "Placeholder":
inputs.append(op.name + ':0')
'''Extract Outputs'''
outputs = [output + ":0" for output in outputs]
print("jrp", outputs)
'''ONNX Graph Generation'''
onnx_graph = tf2onnx.tfonnx.process_tf_graph(sess.graph,
input_names=inputs,
output_names=outputs)
'''Optimizing Grapph for ONNX Formation'''
#opt_graph = tf2onnx.optimizer.optimize_graph(onnx_graph)
'''Make ProtoBuff Model'''
model_proto = onnx_graph.make_model(str(FLAGS.output_path))
#onnx.checker.check_model(model_proto)
'''Store ProtoBuff-file'''
tf2onnx.utils.save_onnx_model("./",
"saved_model",
feed_dict={},
model_proto=model_proto)
print('TF-Graph converted to SavedModel!')
def compute_const_folding_using_tf(g, const_node_values):
"""Find nodes with constant inputs and compute their values using TF"""
if const_node_values is None:
const_node_values = {}
from tf2onnx.tf_loader import tf_session, tf_placeholder # pylint: disable=import-outside-toplevel
ops = g.get_operations()
outputs_to_values = {}
outputs_to_dtypes = {}
for node in ops:
# Load values of constants. Use const_node_values if possible
if node.type in ["Const", "ConstV2"]:
tensor = node.node_def.attr["value"].tensor
if node.name in const_node_values:
tensor.tensor_content = const_node_values[node.name]
outputs_to_values[node.outputs[0].name] = get_tf_tensor_data(
tensor)
outputs_to_dtypes[node.outputs[0].name] = node.outputs[0].dtype
unneeded_outputs = set()
progress = True
while progress:
progress = False
for node in ops:
# Find ops with constant inputs and compute their values
input_names = [i.name for i in node.inputs]
output_names = [i.name for i in node.outputs]
can_fold = node.type not in ['Enter']
can_fold = can_fold and len(input_names) > 0 and all(
inp in outputs_to_values for inp in input_names)
# We can only fold nodes with a single output
can_fold = can_fold and len(
output_names) == 1 and output_names[0] not in outputs_to_values
# Skip if value already computed, used, and discarded
can_fold = can_fold and output_names[0] not in unneeded_outputs
if can_fold:
# Make a mini graph containing just the node to fold
g2 = tf.Graph()
with g2.as_default():
for inp in input_names:
tf_placeholder(outputs_to_dtypes[inp],
name=inp.split(':')[0])
mini_graph_def = g2.as_graph_def()
mini_graph_def.node.append(node.node_def)
g3 = tf.Graph()
with g3.as_default():
feed_dict = {}
for inp in input_names:
feed_dict[inp] = outputs_to_values[inp]
try:
with tf_session() as sess:
tf.import_graph_def(mini_graph_def, name='')
results = sess.run(output_names,
feed_dict=feed_dict)
outputs_to_values[output_names[0]] = results[0]
outputs_to_dtypes[
output_names[0]] = node.outputs[0].dtype
progress = True
except Exception: # pylint: disable=broad-except
logger.debug("Could not fold node %s", node.name)
unneeded_outputs.update(outputs_to_values.keys())
for node in ops:
# Mark values we need to keep
input_names = [i.name for i in node.inputs]
output_names = [i.name for i in node.outputs]
if len(output_names) == 1 and output_names[0] in outputs_to_values:
continue
for i in input_names:
if i in unneeded_outputs:
unneeded_outputs.remove(i)
for node in unneeded_outputs:
# Remove unneeded values to prevent memory usage explosion
if node in outputs_to_values:
del outputs_to_values[node]
del outputs_to_dtypes[node]
for node in ops:
# We don't need the constants any more
if node.type in ["Const", "ConstV2"
] and node.outputs[0].name in outputs_to_values:
del outputs_to_values[node.outputs[0].name]
del outputs_to_dtypes[node.outputs[0].name]
logger.info("Computed %d values for constant folding",
len(outputs_to_values))
return outputs_to_values, outputs_to_dtypes
def run_test_case(self,
func,
feed_dict,
input_names_with_port,
output_names_with_port,
rtol=1e-07,
atol=1e-5,
mtol=None,
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,
premade_placeholders=False,
use_custom_ops=False,
optimize=True):
"""
This function tests all scenarios available through the command line.
The command line always runs the optimizers.
However, they may modify the final graph into something different than the
tested converter implements. Set `optimize=False` to keep the original
set of nodes and helps debugging. However, the same function should
be called with `optimize=True` to test what the user would actually get.
"""
test_tf = not self.config.skip_tf_tests
test_tflite = not self.config.skip_tflite_tests
test_tfjs = not self.config.skip_tfjs_tests
run_tfl_consistency_test = test_tf and test_tflite and self.config.run_tfl_consistency_test
# 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()
if tf_lite is None:
test_tflite = False
g = None
expected, graph_def, initialized_tables = \
self.freeze_and_run_tf(func, feed_dict, output_names_with_port, as_session,
premade_placeholders, large_model, constant_fold)
graph_def_path = os.path.join(
self.test_data_directory,
self._testMethodName + "_after_tf_optimize.pb")
utils.save_protobuf(graph_def_path, graph_def)
self.logger.debug("created file %s", graph_def_path)
if test_tfjs:
tfjs_path = self.convert_to_tfjs(graph_def_path,
output_names_with_port)
if tfjs_path is None:
test_tfjs = False
if test_tflite:
tflite_path = self.convert_to_tflite(graph_def, feed_dict,
output_names_with_port)
test_tflite = tflite_path is not None and self.tflite_has_supported_types(
tflite_path)
if test_tf:
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='')
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,
initialized_tables=initialized_tables,
**process_args)
if optimize:
g = optimizer.optimize_graph(g, catch_errors=False)
actual = self.run_backend(g,
output_names_with_port,
onnx_feed_dict,
large_model,
use_custom_ops=use_custom_ops)
self.assert_results_equal(expected, actual, rtol, atol, mtol,
check_value, check_shape, check_dtype)
self.assert_shapes_correct(g, self.config.allow_missing_shapes,
not self.config.skip_onnx_checker)
if graph_validator:
self.assertTrue(graph_validator(g))
if test_tflite:
tfl_res, tfl_outputs = self.run_tflite(tflite_path, feed_dict)
test_tflite = tfl_res is not None
if test_tflite:
if run_tfl_consistency_test:
self.assert_results_equal(expected, tfl_res, rtol, atol, mtol,
check_value, check_shape,
check_dtype)
tfl_process_args = process_args.copy()
if 'inputs_as_nchw' in tfl_process_args:
nchw_inps_with_port = tfl_process_args['inputs_as_nchw']
tfl_process_args['inputs_as_nchw'] = [
i.split(':')[0] for i in nchw_inps_with_port
]
input_names_without_port = [
inp.split(':')[0] for inp in feed_dict.keys()
]
g = process_tf_graph(None,
opset=self.config.opset,
input_names=input_names_without_port,
output_names=tfl_outputs,
target=self.config.target,
tflite_path=tflite_path,
**tfl_process_args)
if optimize:
g = optimizer.optimize_graph(g)
onnx_feed_dict_without_port = {
k.split(':')[0]: v
for k, v in onnx_feed_dict.items()
}
onnx_tfl_res = self.run_backend(g,
tfl_outputs,
onnx_feed_dict_without_port,
postfix="_from_tflite",
use_custom_ops=use_custom_ops)
self.assert_results_equal(tfl_res, onnx_tfl_res, rtol, atol, mtol,
check_value, check_shape, check_dtype)
self.assert_shapes_correct(g, self.config.allow_missing_shapes,
not self.config.skip_onnx_checker)
if graph_validator:
self.assertTrue(graph_validator(g))
if test_tfjs:
try:
tfjs_res = run_tfjs(tfjs_path, feed_dict)
except RuntimeError as e:
ignored_errors = [
"is not yet supported",
"Operands could not be broadcast together",
"unknown dtype null", "must be [NaN",
"Cannot read property 'name' of undefined",
"Either strides or dilations must be 1", "does not support"
]
if any(err in str(e) for err in ignored_errors):
test_tfjs = False
else:
raise e
if test_tfjs:
g = process_tf_graph(None,
opset=self.config.opset,
input_names=list(feed_dict.keys()),
output_names=None,
target=self.config.target,
tfjs_path=tfjs_path,
**process_args)
g = optimizer.optimize_graph(g)
onnx_tfjs_res = self.run_backend(g,
None,
onnx_feed_dict,
large_model,
postfix="_from_tfjs",
use_custom_ops=use_custom_ops)
self.assert_results_equal(tfjs_res,
onnx_tfjs_res,
rtol,
atol,
mtol,
check_value,
check_shape,
check_dtype=False)
self.assert_shapes_correct(g, self.config.allow_missing_shapes,
not self.config.skip_onnx_checker)
if graph_validator:
self.assertTrue(graph_validator(g))
if g is None:
raise unittest.SkipTest("tf, tflite, and tfjs marked to skip")
return g
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 run_test(self,
name,
backend="onnxruntime",
onnx_file=None,
opset=None,
extra_opset=None,
perf=None):
"""Run complete test against backend."""
self.perf = perf
# get the model
if self.url:
_, dir_name = self.download_model()
logger.info("Downloaded to %s", dir_name)
model_path = os.path.join(
dir_name, self.local) if self.local != "." else dir_name
else:
model_path = self.local
logger.info("Load model from %s", model_path)
input_names = list(self.input_names.keys())
initialized_tables = {}
outputs = self.output_names
tflite_path = None
to_rename = None
if self.model_type in ["checkpoint"]:
graph_def, input_names, outputs = tf_loader.from_checkpoint(
model_path, input_names, outputs)
elif self.model_type in ["saved_model"]:
loaded = tf_loader.from_saved_model(
model_path,
None,
None,
self.tag,
self.signatures,
self.concrete_function,
self.large_model,
return_concrete_func=not self.run_tf_frozen,
return_initialized_tables=True,
return_tensors_to_rename=True)
if not self.run_tf_frozen:
# Must maintain ref to imported since concrete_func uses weak refs
# pylint: disable=unused-variable
graph_def, input_names, outputs, concrete_func, imported, initialized_tables, to_rename = loaded
else:
graph_def, input_names, outputs, initialized_tables, to_rename = loaded
elif self.model_type in ["keras"]:
graph_def, input_names, outputs = tf_loader.from_keras(
model_path, input_names, outputs)
elif self.model_type in ["tflite"]:
tflite_path = model_path
graph_def = None
else:
graph_def, input_names, outputs = tf_loader.from_graphdef(
model_path, input_names, outputs)
if utils.is_debug_mode():
utils.save_protobuf(
os.path.join(TEMP_DIR, name + "_after_tf_optimize.pb"),
graph_def)
if tflite_path is not None:
inputs = {}
for k in input_names:
v = self.input_names[k]
inputs[k] = self.make_input(v)
interpreter = tf.lite.Interpreter(tflite_path)
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
input_name_to_index = {
n['name'].split(':')[0]: n['index']
for n in input_details
}
for k, v in inputs.items():
interpreter.resize_tensor_input(input_name_to_index[k],
v.shape)
interpreter.allocate_tensors()
def run_tflite():
for k, v in inputs.items():
interpreter.set_tensor(input_name_to_index[k], v)
interpreter.invoke()
result = [
interpreter.get_tensor(output['index'])
for output in output_details
]
return result
tf_results = run_tflite()
if self.perf:
logger.info("Running TFLite perf")
n = 0
start = time.time()
stop = start + PERF_TIME
while time.time() < stop:
for _ in range(PERF_STEP):
_ = run_tflite()
n += PERF_STEP
self.tf_runtime = 1000 * (time.time() - start) / n
logger.info("TFLite perf {:.2f}ms/inference, n={}".format(
self.tf_runtime, n))
logger.info("TFLite OK")
if not self.run_tf_frozen:
inputs = {}
for k in input_names:
v = self.input_names[k]
inputs[k.split(":")[0]] = tf.constant(self.make_input(v))
tf_func = tf.function(concrete_func)
logger.info("Running TF")
tf_results_d = tf_func(**inputs)
# If there is only a single output a dict might not be returned
if isinstance(tf_results_d, tf.Tensor):
tf_results = [tf_results_d]
else:
tf_results = [
tf_results_d[k] for k in sorted(tf_results_d.keys())
]
tf_results = [tf_res.numpy() for tf_res in tf_results]
if self.perf:
logger.info("Running TF perf")
n = 0
start = time.time()
stop = start + PERF_TIME
if self.tf_profile is not None:
tf.profiler.experimental.start(self.tf_profile)
while time.time() < stop:
for _ in range(PERF_STEP):
_ = concrete_func(**inputs)
n += PERF_STEP
if self.tf_profile is not None:
tf.profiler.experimental.stop()
self.tf_runtime = 1000 * (time.time() - start) / n
logger.info("TF perf {:.2f}ms/inference, n={}".format(
self.tf_runtime, n))
logger.info("TensorFlow OK")
shape_override = {}
const_node_values = None
tf_graph = None
if graph_def is not None:
inputs = {}
tf_reset_default_graph()
with tf.Graph().as_default() as tf_graph:
from tf2onnx.tf_utils import compress_graph_def
if self.large_model:
const_node_values = compress_graph_def(graph_def)
tf.import_graph_def(graph_def, name='')
with tf_session(graph=tf_graph) as sess:
# create the input data
for k in input_names:
v = self.input_names[k]
t = sess.graph.get_tensor_by_name(k)
expected_dtype = tf.as_dtype(t.dtype).name
if isinstance(v, six.text_type) and v.startswith("np."):
np_value = eval(v) # pylint: disable=eval-used
if expected_dtype != np_value.dtype:
logger.warning(
"dtype mismatch for input %s: expected=%s, actual=%s",
k, expected_dtype, np_value.dtype)
inputs[k] = np_value.astype(expected_dtype)
else:
if expected_dtype == "string":
inputs[k] = self.make_input(v).astype(
np.str).astype(np.object)
else:
inputs[k] = self.make_input(v).astype(
expected_dtype)
if self.force_input_shape:
for k, v in inputs.items():
shape_override[k] = list(v.shape)
# run the model with tensorflow
if self.skip_tensorflow:
logger.info("TensorFlow SKIPPED")
elif self.run_tf_frozen:
if self.tf_profile is not None:
tf.profiler.experimental.start(self.tf_profile)
tf_results = self.run_tensorflow(sess, inputs)
if self.tf_profile is not None:
tf.profiler.experimental.stop()
logger.info("TensorFlow OK")
tf_graph = sess.graph
model_proto = None
if self.skip_conversion:
if self.large_model:
external_tensor_storage = ExternalTensorStorage()
model_proto = utils.model_proto_from_zip(
self.converted_model, external_tensor_storage)
else:
external_tensor_storage = None
model_proto = utils.model_proto_from_file(self.converted_model)
logger.info("ONNX loaded from file")
else:
try:
# convert model to onnx
onnx_graph = self.to_onnx(
tf_graph,
opset=opset,
extra_opset=extra_opset,
shape_override=shape_override,
input_names=inputs.keys(),
const_node_values=const_node_values,
initialized_tables=initialized_tables,
tflite_path=tflite_path,
tensors_to_rename=to_rename)
onnx_graph = optimizer.optimize_graph(onnx_graph)
print("ONNX", onnx_graph.dump_node_statistics())
external_tensor_storage = ExternalTensorStorage(
) if self.large_model else None
model_proto = onnx_graph.make_model(
"converted from tf2onnx",
external_tensor_storage=external_tensor_storage)
logger.info("To_ONNX, OK")
if onnx_file:
self.create_onnx_file(name, model_proto, inputs, onnx_file,
external_tensor_storage)
if self.converted_model:
if self.large_model:
utils.save_onnx_zip(self.converted_model, model_proto,
external_tensor_storage)
else:
utils.save_protobuf(self.converted_model, model_proto)
logger.info("Created %s", self.converted_model)
except Exception:
logger.error("To_ONNX FAIL", exc_info=1)
return False
try:
onnx_results = None
if backend == "onnxruntime":
if to_rename is None:
struc_outputs = self.output_names
else:
struc_outputs = [
to_rename.get(k, k) for k in self.output_names
]
onnx_results = self.run_onnxruntime(name, model_proto, inputs,
struc_outputs,
external_tensor_storage)
else:
raise ValueError("unknown backend")
logger.info("Run_ONNX OK")
try:
if self.skip_tensorflow:
logger.info("Results: skipped tensorflow")
else:
if self.check_only_shape:
for tf_res, onnx_res in zip(tf_results, onnx_results):
np.testing.assert_array_equal(
tf_res.shape, onnx_res.shape)
else:
for tf_res, onnx_res in zip(tf_results, onnx_results):
good_cnt = np.count_nonzero(
np.isclose(tf_res,
onnx_res,
rtol=self.rtol,
atol=self.atol))
bad_cnt = tf_res.size - good_cnt
if bad_cnt > self.ptol / 100 * tf_res.size:
# Prints a nice error message with stats
np.testing.assert_allclose(tf_res,
onnx_res,
rtol=self.rtol,
atol=self.atol)
logger.info("Results: OK")
return True
except Exception:
logger.error("Results", exc_info=1)
except Exception:
logger.error("Run_ONNX FAIL", exc_info=1)
return False
def main():
args = get_args()
logging.basicConfig(level=logging.get_verbosity_level(args.verbose))
if args.debug:
utils.set_debug_mode(True)
logger = logging.getLogger(constants.TF2ONNX_PACKAGE_NAME)
extra_opset = args.extra_opset or []
custom_ops = {}
if args.custom_ops:
# default custom ops for tensorflow-onnx are in the "tf" namespace
custom_ops = {
op: (default_custom_op_handler, [])
for op in args.custom_ops.split(",")
}
extra_opset.append(constants.TENSORFLOW_OPSET)
# get the frozen tensorflow model from graphdef, checkpoint or saved_model.
if args.graphdef:
graph_def, inputs, outputs = tf_loader.from_graphdef(
args.graphdef, args.inputs, args.outputs)
model_path = args.graphdef
if args.checkpoint:
graph_def, inputs, outputs = tf_loader.from_checkpoint(
args.checkpoint, args.inputs, args.outputs)
model_path = args.checkpoint
if args.saved_model:
graph_def, inputs, outputs = tf_loader.from_saved_model(
args.saved_model, args.inputs, args.outputs, args.signature_def)
model_path = args.saved_model
if args.keras:
graph_def, inputs, outputs = tf_loader.from_keras(
args.keras, args.inputs, args.outputs)
model_path = args.keras
if args.verbose:
logger.info("inputs: %s", inputs)
logger.info("outputs: %s", outputs)
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
with tf_loader.tf_session(graph=tf_graph):
g = process_tf_graph(tf_graph,
continue_on_error=args.continue_on_error,
target=args.target,
opset=args.opset,
custom_op_handlers=custom_ops,
extra_opset=extra_opset,
shape_override=args.shape_override,
input_names=inputs,
output_names=outputs,
inputs_as_nchw=args.inputs_as_nchw)
onnx_graph = optimizer.optimize_graph(g)
model_proto = onnx_graph.make_model("converted from {}".format(model_path))
# write onnx graph
logger.info("")
logger.info("Successfully converted TensorFlow model %s to ONNX",
model_path)
if args.output:
utils.save_protobuf(args.output, model_proto)
logger.info("ONNX model is saved at %s", args.output)
else:
logger.info(
"To export ONNX model to file, please run with `--output` option")
if (output_name not in nodes_inputs) and (0 < len(n.input)):
graph_output_names.add(output_name)
logging.info('Model Inputs: %s', str(list(graph_input_names)))
logging.info('Model Outputs: %s', str(list(graph_output_names)))
graph_def, inputs, outputs = tf_loader.from_graphdef(
model_path=args.in_file,
input_names=list(graph_input_names),
output_names=list(graph_output_names))
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
if 160 <= TF2ONNX_VERSION:
with tf_loader.tf_session(graph=tf_graph):
onnx_graph = tf2onnx.tfonnx.process_tf_graph(tf_graph=tf_graph,
input_names=inputs,
output_names=outputs,
opset=9)
else:
with tf.Session(graph=tf_graph):
onnx_graph = tf2onnx.tfonnx.process_tf_graph(tf_graph=tf_graph,
input_names=inputs,
output_names=outputs,
opset=9)
# Optimize with tf2onnx.optimizer
onnx_graph = tf2onnx.optimizer.optimize_graph(onnx_graph)
model_proto = onnx_graph.make_model(model_name)