def main():
args = get_args()
opset = tf2onnx.utils.find_opset(args.opset)
print("using tensorflow={}, onnx={}, opset={}, tfonnx={}/{}".format(
tf.__version__, onnx.__version__, opset, tf2onnx.__version__,
tf2onnx.version.git_version[:6]))
# override unknown dimensions from -1 to 1 (aka batchsize 1) since not every runtime does
# support unknown dimensions.
tf2onnx.utils.ONNX_UNKNOWN_DIMENSION = args.unknown_dim
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 = [helper.make_opsetid(_TENSORFLOW_DOMAIN, 1)]
else:
custom_ops = {}
extra_opset = None
graph_def = tf.GraphDef()
with tf.gfile.GFile(args.input, 'rb') as f:
graph_def.ParseFromString(f.read())
# todo: consider to enable const folding by default?
graph_def = tf_optimize(args.inputs, args.outputs, graph_def,
args.fold_const)
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
with tf.Session(graph=tf_graph):
g = process_tf_graph(tf_graph,
continue_on_error=args.continue_on_error,
verbose=args.verbose,
target=args.target,
opset=args.opset,
custom_op_handlers=custom_ops,
extra_opset=extra_opset,
shape_override=args.shape_override,
input_names=args.inputs,
output_names=args.outputs,
inputs_as_nchw=args.inputs_as_nchw)
new_model_proto = GraphUtil.opt_transposes_with_graph(
g,
"converted from {}".format(args.input),
optimize=not args.continue_on_error)
if new_model_proto:
model_proto = new_model_proto
else:
print("NON-CRITICAL, optimizers are not applied successfully")
# write onnx graph
if args.output:
with open(args.output, "wb") as f:
f.write(model_proto.SerializeToString())
print("\nComplete successfully, the onnx model is generated at " +
args.output)
def main():
args = get_args()
print("using tensorflow={}, onnx={}".format(tf.__version__,
onnx.__version__))
# override unknown dimensions from -1 to 1 (aka batchsize 1) since not every runtime does
# support unknown dimensions.
tf2onnx.utils.ONNX_UNKNOWN_DIMENSION = args.unknown_dim
graph_def = tf.GraphDef()
with tf.gfile.FastGFile(args.input, 'rb') as f:
graph_def.ParseFromString(f.read())
graph_def = tf_optimize(None, args.inputs, args.outputs, graph_def)
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
with tf.Session(graph=tf_graph) as sess:
g = process_tf_graph(tf_graph,
continue_on_error=args.continue_on_error,
verbose=args.verbose,
target=args.target,
opset=args.opset)
model_proto = g.make_model("converted from {}".format(args.input),
args.inputs, args.outputs)
# write onnx graph
if args.output:
with open(args.output, "wb") as f:
f.write(model_proto.SerializeToString())
def convert_to_onnx(model_filepath, onnx_filepath, output_node_names):
"""
Convert the model to an ONNX file, which can in turn be used for TensorRT inference
Arguments:
model_filepath: the path to the frozen .pb file
onnx_filepath: the path where the ONNX file should be saved
output_node_names: list of output node names
"""
# tf2onnx expects the node names in the format "input/output_node_name:port_id".
# Hence, we should provide the port ID before conversion.
input_node_names = [kInputName + ":0"]
output_node_names = list(map(lambda x: x + ":0", output_node_names))
# Use in-built function from tf2onnx to import the graph and optimize for conversion
graph_def, inputs, outputs = loader.from_graphdef(model_filepath,
input_node_names,
output_node_names)
graph_def = tf_optimize(input_node_names, output_node_names, graph_def,
False)
with tf.Graph().as_default() as default_graph:
tf.import_graph_def(graph_def, name='')
# Convert to ONNX
with tf.Session(graph=default_graph):
onnx_graph = process_tf_graph(default_graph,
opset=8,
input_names=inputs,
output_names=outputs)
onnx_graph = optimizer.optimize_graph(onnx_graph)
onnx_model = onnx_graph.make_model("segmentation_onnx_model")
# Save the ONNX model to disk
utils.save_protobuf(onnx_filepath, onnx_model)
def convert_frozen_to_onnx(
settings: SerializationSettings, frozen_graph_def: tf.GraphDef
) -> Any:
# This is basically https://github.com/onnx/tensorflow-onnx/blob/master/tf2onnx/convert.py
inputs = _get_input_node_names(frozen_graph_def)
outputs = _get_output_node_names(frozen_graph_def)
logger.info(f"onnx export - inputs:{inputs} outputs:{outputs}")
frozen_graph_def = tf_optimize(
inputs, outputs, frozen_graph_def, fold_constant=True
)
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(frozen_graph_def, name="")
with tf.Session(graph=tf_graph):
g = process_tf_graph(
tf_graph,
input_names=inputs,
output_names=outputs,
opset=settings.onnx_opset,
)
onnx_graph = optimizer.optimize_graph(g)
model_proto = onnx_graph.make_model(settings.brain_name)
return model_proto
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 = loader.from_graphdef(args.graphdef, args.inputs, args.outputs)
model_path = args.graphdef
if args.checkpoint:
graph_def, inputs, outputs = loader.from_checkpoint(args.checkpoint, args.inputs, args.outputs)
model_path = args.checkpoint
if args.saved_model:
graph_def, inputs, outputs = loader.from_saved_model(
args.saved_model, args.inputs, args.outputs, args.signature_def)
model_path = args.saved_model
if args.verbose:
logger.info("inputs: %s", inputs)
logger.info("outputs: %s", outputs)
# todo: consider to enable const folding by default?
graph_def = tf_optimize(inputs, outputs, graph_def, args.fold_const)
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
with 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")
def run_test_case(self, feed_dict, input_names_with_port, output_names_with_port, rtol=1e-07,
convert_var_to_const=True, constant_fold=True, check_value=True, check_shape=False,
check_dtype=False, process_args=None, onnx_feed_dict=None):
# 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
graph_def = None
if convert_var_to_const:
with tf.Session() as sess:
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='')
with tf.Session() as sess:
variables_lib.global_variables_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)
if self.debug_mode():
model_path = os.path.join(type(self).TMPPATH, self._testMethodName + "_original.pb")
with open(model_path, "wb") as f:
f.write(sess.graph_def.SerializeToString())
self.log.debug("created file %s", model_path)
graph_def = tf_optimize(input_names_with_port, output_names_with_port,
sess.graph_def, constant_fold)
if self.debug_mode() and constant_fold:
model_path = os.path.join(type(self).TMPPATH, self._testMethodName + "_after_tf_optimize.pb")
with open(model_path, "wb") as f:
f.write(graph_def.SerializeToString())
self.log.debug("created file %s", model_path)
tf.reset_default_graph()
tf.import_graph_def(graph_def, name='')
with tf.Session() as sess:
g = process_tf_graph(sess.graph, opset=type(self).OPSET, output_names=output_names_with_port,
**process_args)
actual = self._run_backend(g, output_names_with_port, onnx_feed_dict)
for expected_val, actual_val in zip(expected, actual):
if check_value:
self.assertAllClose(expected_val, actual_val, rtol=rtol, atol=0.)
if check_dtype:
self.assertEqual(expected_val.dtype, actual_val.dtype)
if check_shape:
self.assertEqual(expected_val.shape, actual_val.shape)
def convert_tensorflow(frozen_graph_def,
name=None,
input_names=None,
output_names=None,
doc_string='',
target_opset=None,
channel_first_inputs=None,
debug_mode=False,
custom_op_conversions=None):
"""
convert a frozen tensorflow graph def into a ONNX model proto, just like how keras does.
:param frozen_graph_def: the frozen tensorflow graph
:param name: the converted onnx model internal name
:param input_names: the inputs name list of the model
:param output_names: the output name list of the model
:param doc_string: doc string
:param target_opset: the targeted onnx model opset
:param channel_first_inputs: A list of channel first input (not supported yet)
:param debug_mode: will enable the log and try to convert as much as possible on conversion
:return an ONNX ModelProto
"""
set_logger_level(logging.DEBUG if debug_mode else logging.INFO)
if target_opset is None:
target_opset = get_opset_number_from_onnx()
if not doc_string:
doc_string = "converted from {}".format(name)
graph_def = tfonnx.tf_optimize(input_names, output_names, frozen_graph_def,
True)
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
if get_tensorboard_writer() is not None:
get_tensorboard_writer().add_graph(tf_graph)
custom_op_handlers = tf2onnx_builtin_conversion(target_opset)
if custom_op_conversions:
custom_op_handlers.update(custom_op_conversions)
with tf.Session(graph=tf_graph):
if not input_names:
input_nodes = list(_collect_input_nodes(tf_graph, output_names)[0])
input_names = [nd_.outputs[0].name for nd_ in input_nodes]
g = tfonnx.process_tf_graph(tf_graph,
continue_on_error=debug_mode,
opset=target_opset,
custom_op_handlers=custom_op_handlers,
inputs_as_nchw=channel_first_inputs,
output_names=output_names,
input_names=input_names)
onnx_graph = tf2onnx.optimizer.optimize_graph(g)
model_proto = onnx_graph.make_model(doc_string)
return model_proto
def main():
args = get_args()
opset = tf2onnx.utils.find_opset(args.opset)
print("using tensorflow={}, onnx={}, opset={}".format(
tf.__version__, onnx.__version__, opset))
# override unknown dimensions from -1 to 1 (aka batchsize 1) since not every runtime does
# support unknown dimensions.
tf2onnx.utils.ONNX_UNKNOWN_DIMENSION = args.unknown_dim
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 = [helper.make_opsetid(_TENSORFLOW_DOMAIN, 1)]
else:
custom_ops = {}
extra_opset = None
graph_def = tf.GraphDef()
with tf.gfile.FastGFile(args.input, 'rb') as f:
graph_def.ParseFromString(f.read())
# todo: consider to enable const folding by default?
graph_def = tf_optimize(args.inputs, args.outputs, graph_def,
args.fold_const)
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
with tf.Session(graph=tf_graph):
g = process_tf_graph(tf_graph,
continue_on_error=args.continue_on_error,
verbose=args.verbose,
target=args.target,
opset=args.opset,
custom_op_handlers=custom_ops,
extra_opset=extra_opset,
shape_override=args.shape_override)
optimizer = TransposeOptimizer(g, args.verbose is not None)
optimizer.optimize()
model_proto = g.make_model("converted from {}".format(args.input),
args.outputs,
optimize=not args.continue_on_error)
# write onnx graph
if args.output:
with open(args.output, "wb") as f:
f.write(model_proto.SerializeToString())
def convert_onnx(sess, graph_def, input_path, inputs_op, outputs_op):
graphdef = input_path
if inputs_op:
inputs_op, shape_override = utils.split_nodename_and_shape(inputs_op)
if outputs_op:
outputs_op = outputs_op.split(",")
logging.basicConfig(level=logging.get_verbosity_level(True))
utils.set_debug_mode(True)
logger = logging.getLogger(constants.TF2ONNX_PACKAGE_NAME)
graph_def, inputs_op, outputs_op = from_graphdef(sess, graph_def, graphdef,
inputs_op, outputs_op)
model_path = graphdef
graph_def = tf_optimize(inputs_op, outputs_op, graph_def, True)
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
with tf.Session(graph=tf_graph):
g = process_tf_graph(tf_graph,
continue_on_error=False,
target=",".join(constants.DEFAULT_TARGET),
opset=10,
custom_op_handlers=None,
extra_opset=None,
shape_override=None,
input_names=inputs_op,
output_names=outputs_op,
inputs_as_nchw=None)
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:
output_path = input_path.replace(".pb", ".onnx")
utils.save_protobuf(output_path, model_proto)
logger.info("ONNX model is saved at %s", output_path)
def convert_frozen_to_onnx(settings: SerializationSettings,
frozen_graph_def: tf.GraphDef) -> Any:
# This is basically https://github.com/onnx/tensorflow-onnx/blob/master/tf2onnx/convert.py
# Some constants in the graph need to be read by the inference system.
# These aren't used by the model anywhere, so trying to make sure they propagate
# through conversion and import is a losing battle. Instead, save them now,
# so that we can add them back later.
constant_values = {}
for n in frozen_graph_def.node:
if n.name in MODEL_CONSTANTS:
val = n.attr["value"].tensor.int_val[0]
constant_values[n.name] = val
inputs = _get_input_node_names(frozen_graph_def)
outputs = _get_output_node_names(frozen_graph_def)
logger.info(f"onnx export - inputs:{inputs} outputs:{outputs}")
frozen_graph_def = tf_optimize(inputs,
outputs,
frozen_graph_def,
fold_constant=True)
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(frozen_graph_def, name="")
with tf.Session(graph=tf_graph):
g = process_tf_graph(
tf_graph,
input_names=inputs,
output_names=outputs,
opset=settings.onnx_opset,
)
onnx_graph = optimizer.optimize_graph(g)
model_proto = onnx_graph.make_model(settings.brain_name)
# Save the constant values back the graph initializer.
# This will ensure the importer gets them as global constants.
constant_nodes = []
for k, v in constant_values.items():
constant_node = _make_onnx_node_for_constant(k, v)
constant_nodes.append(constant_node)
model_proto.graph.initializer.extend(constant_nodes)
return model_proto
def _convert(self):
logger.info("Converting...")
self._restore_session_from_model_file()
tf.reset_default_graph()
graph_def = tf_optimize(self.original_input_names, self.original_output_names,
self.tf_sess.graph_def)
tf.import_graph_def(graph_def, name="")
onnx_graph = process_tf_graph(tf.get_default_graph(), opset=self.conversion_config.onnx_opset,
input_names=self.original_input_names, output_names=self.original_output_names)
try:
opt_graph = optimizer.optimize_graph(onnx_graph)
except Exception as e:
opt_graph = None
logger.warning("Failed to optimize ONNX graph, original un-optimized graph will be used, e = {}".format(e))
onnx_graph = opt_graph if opt_graph is not None else onnx_graph
model_proto = onnx_graph.make_model("onnx-proto")
with open(self.conversion_config.onnx_model_path, "wb") as f:
f.write(model_proto.SerializeToString())
def _run_test_case(self, input_names_with_port, output_names_with_port):
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
graph_def = tf_optimize(input_names_with_port, output_names_with_port,
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_tf2onnx(model,
output,
inputs,
outputs,
signature_def=None,
opset=7):
import tensorflow as tf
from tf2onnx.tfonnx import process_tf_graph, tf_optimize
from tf2onnx import constants, loader, logging, utils, optimizer
logger = logging.getLogger(constants.TF2ONNX_PACKAGE_NAME)
if "pb" in model:
graph_def, inputs, outputs = loader.from_graphdef(
model, inputs, outputs)
elif "meta" in model:
graph_def, inputs, outputs = loader.from_checkpoint(
model, inputs, outputs)
elif "saved_model" in model:
graph_def, inputs, outputs = loader.from_saved_model(
model, inputs, outputs, signature_def)
graph_def = tf_optimize(inputs, outputs, graph_def, None)
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
with tf.Session(graph=tf_graph):
g = process_tf_graph(tf_graph,
opset=opset,
input_names=inputs,
output_names=outputs)
onnx_graph = optimizer.optimize_graph(g)
model_proto = onnx_graph.make_model("converted from {}".format(model))
# write onnx graph
logger.info("")
logger.info("Successfully converted TensorFlow model %s to ONNX", model)
utils.save_protobuf(output, model_proto)
logger.info("ONNX model is saved at %s", output)
tf.import_graph_def(graph_def, name='')
# [array([[[-0.05846359 -0.06566401 0.02254938 -0.26033643 -0.07923548]],
# [[ 0.04879569 0.04215769 -0.06720451 -0.60583305 0.06223793]],
# [[-0.05626901 -0.06627436 0.00422506 -0.5533649 -0.0767431 ]]], dtype=float32)]
with tf.Session() as sess:
# output_dict: get tensor by output name
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={"input_1:0": x_val})
# tf optimize
graph_def = tf_optimize(input_names_with_port,
output_names_with_port,
sess.graph_def,
fold_constant=True)
tf.reset_default_graph()
tf.import_graph_def(graph_def, name='')
# convert to onnx
with tf.Session() as sess:
g = process_tf_graph(sess.graph, output_names=output_names_with_port)
g = optimizer.optimize_graph(g)
model_proto = g.make_model("lstm")
utils.save_onnx_model("./models",
"lstm",
feed_dict={"input_1:0": input},
model_proto=model_proto)
def run_test_case(self,
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=False,
check_dtype=True,
process_args=None,
onnx_feed_dict=None,
graph_validator=None):
# 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
graph_def = None
if convert_var_to_const:
with tf.Session() as sess:
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='')
with tf.Session() as sess:
variables_lib.global_variables_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)
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 + "_original.pb")
utils.save_protobuf(model_path, sess.graph_def)
self.log.debug("created file %s", model_path)
graph_def = tf_optimize(input_names_with_port, output_names_with_port,
sess.graph_def, constant_fold)
if self.config.is_debug_mode and 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.log.debug("created file %s", model_path)
tf.reset_default_graph()
tf.import_graph_def(graph_def, name='')
with tf.Session() as sess:
g = process_tf_graph(sess.graph,
opset=self.config.opset,
output_names=output_names_with_port,
target=self.config.target,
**process_args)
g = optimizer.optimize_graph(g)
actual = self._run_backend(g, output_names_with_port,
onnx_feed_dict)
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)
if check_shape:
self.assertEqual(expected_val.shape, actual_val.shape)
if graph_validator:
self.assertTrue(graph_validator(g))
return g
def pb_to_onnx(
inputs: List[Union[str, tf_compat.Tensor]],
outputs: List[Union[str, tf_compat.Tensor]],
pb_path: str,
onnx_path: str,
opset: int = default_onnx_opset(),
custom_op_handlers=None,
extra_opset=None,
shape_override: Dict[str, List] = None,
):
"""
Export an ONNX format for the graph from PB format.
Should not be called within an active graph or session.
:param inputs: the inputs the graph should be created for,
can be either a list of names or a list of tensors
:param outputs: the outputs the graph should be created for,
can be either a list of names or a list of tensors
:param pb_path: path to the existing PB file
:param onnx_path: path to the output ONNX file
:param opset: ONNX opset
:param custom_op_handlers: dictionary of custom op handlers
:param extra_opset: list of extra opset's
:param shape_override: new shape to override
"""
try:
from tf2onnx import constants, optimizer, utils
from tf2onnx.tfonnx import process_tf_graph, tf_optimize
except ModuleNotFoundError:
raise ModuleNotFoundError(
"tf2onnx must be installed on the system before using export_onnx"
)
try:
from tf2onnx import tf_loader as loader
except Exception:
from tf2onnx import loader
pb_path = clean_path(pb_path)
if not os.path.exists(pb_path):
raise FileNotFoundError(
("no pb file for the model found at {}").format(pb_path)
)
inputs = [inp if isinstance(inp, str) else inp.name for inp in inputs]
outputs = [out if isinstance(out, str) else out.name for out in outputs]
graph_def, inputs, outputs = loader.from_graphdef(pb_path, inputs, outputs)
graph_def = tf_optimize(inputs, outputs, graph_def, fold_constant=True)
with tf_compat.Graph().as_default() as tf_graph:
tf_compat.import_graph_def(graph_def, name="")
with tf_compat.Session(graph=tf_graph):
graph = process_tf_graph(
tf_graph,
continue_on_error=False,
target=",".join(constants.DEFAULT_TARGET),
opset=opset,
custom_op_handlers=custom_op_handlers,
extra_opset=extra_opset,
shape_override=shape_override,
input_names=inputs,
output_names=outputs,
)
onnx_graph = optimizer.optimize_graph(graph)
model_proto = onnx_graph.make_model("converted from {}".format(pb_path))
onnx_path = clean_path(onnx_path)
create_parent_dirs(onnx_path)
utils.save_protobuf(onnx_path, model_proto)
def main():
args = get_args()
logging.basicConfig(level=logging.get_verbosity_level(args.verbose))
if args.debug:
utils.set_debug_mode(True)
# override unknown dimensions from -1 to 1 (aka batchsize 1) since not every runtime does
# support unknown dimensions.
utils.ONNX_UNKNOWN_DIMENSION = args.unknown_dim
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 = loader.from_graphdef(
args.graphdef, args.inputs, args.outputs)
model_path = args.graphdef
if args.checkpoint:
graph_def, inputs, outputs = loader.from_checkpoint(
args.checkpoint, args.inputs, args.outputs)
model_path = args.checkpoint
if args.saved_model:
graph_def, inputs, outputs = loader.from_saved_model(
args.saved_model, args.inputs, args.outputs)
model_path = args.saved_model
# todo: consider to enable const folding by default?
graph_def = tf_optimize(inputs, outputs, graph_def, args.fold_const)
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
with 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)
model_proto = g.make_model("converted from {}".format(model_path))
new_model_proto = GraphUtil.optimize_model_proto(model_proto)
if new_model_proto:
model_proto = new_model_proto
else:
print("NON-CRITICAL, optimizers are not applied successfully")
# write onnx graph
if args.output:
utils.save_protobuf(args.output, model_proto)
print("\nComplete successfully, the onnx model is generated at " +
args.output)
def main():
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
args = get_args()
print("using tensorflow={}, onnx={}".format(tf.__version__, onnx.__version__))
# override unknown dimensions from -1 to 1 (aka batchsize 1) since not every runtime does
# support unknown dimensions.
tf2onnx.utils.ONNX_UNKNOWN_DIMENSION = args.unknown_dim
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 = [helper.make_opsetid(_TENSORFLOW_DOMAIN, 1)]
else:
custom_ops = {}
extra_opset = None
graph_def = tf.GraphDef()
with tf.gfile.FastGFile(args.input, 'rb') as f:
graph_def.ParseFromString(f.read())
print("args.inputs:", args.inputs)
print("args.outputs:", args.outputs)
print("args.middle_inputs:", args.middle_inputs)
print("args.middle_outputs:", args.middle_outputs)
if args.middle_outputs:
graph_def = tf_optimize(None, args.inputs, args.middle_outputs, graph_def)
else:
graph_def = tf_optimize(None, args.inputs, args.outputs, graph_def)
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
with tf.Session(graph=tf_graph) as sess:
g = process_tf_graph(tf_graph,
continue_on_error=args.continue_on_error,
verbose=args.verbose,
target=args.target,
opset=args.opset,
custom_op_handlers=custom_ops,
extra_opset=extra_opset,
shape_override=args.shape_override,
inputs=args.inputs,
middle_inputs=args.middle_inputs)
final_inputs = args.inputs
final_outputs = args.outputs
if args.middle_outputs:
final_outputs = args.middle_outputs
if args.middle_inputs:
final_inputs = args.middle_inputs
model_proto = g.make_model(
"converted from {}".format(args.input), final_inputs, final_outputs,
optimize=not args.continue_on_error)
# write onnx graph
if args.output:
with open(args.output, "wb") as f:
f.write(model_proto.SerializeToString())
def convert2onnx(self):
inputs = [
'%s:0' % (layer.name) for layer in self.lwnn_model
if layer.op == 'Input'
]
if ('output_node' in self.kwargs):
outputs = ['%s:0' % (o) for o in self.kwargs['output_node']]
else:
outputs = [
'%s:0' % (layer.name) for layer in self.lwnn_model
if len(self.get_consumers(layer)) == 0
]
custom_ops = {}
extra_opset = []
graph_def = tf_optimize(inputs, outputs, self.graph_def, True)
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
with tf.Session(graph=tf_graph):
g = process_tf_graph(tf_graph,
continue_on_error=False,
target=None,
opset=None,
custom_op_handlers=custom_ops,
extra_opset=extra_opset,
shape_override=None,
input_names=inputs,
output_names=outputs,
inputs_as_nchw=None)
onnx_graph = tf2onnx.optimizer.optimize_graph(g)
model = onnx_graph.make_model(self.name)
inps = []
for inp in model.graph.input:
if (inp.name in inputs):
shape = [
int(dim.dim_value)
for dim in inp.type.tensor_type.shape.dim
]
if (len(shape) == 0):
layer = self.get_layers(self.LN(inp.name))
x = onnx.helper.make_tensor_value_info(
inp.name, inp.type.tensor_type.elem_type, layer.shape)
inps.append(x)
else:
inps.append(inp)
outs = []
for out in model.graph.output:
if (out.name in inputs):
shape = [
int(dim.dim_value)
for dim in out.type.tensor_type.shape.dim
]
if (len(shape) == 0):
layer = self.get_layers(self.LN(out.name))[0]
x = onnx.helper.make_tensor_value_info(
out.name, out.type.tensor_type.elem_type, layer.shape)
inps.append(x)
else:
inps.append(out)
del model.graph.input[:]
model.graph.input.extend(inps)
del model.graph.output[:]
model.graph.output.extend(outs)
return model
