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 main(input_file, output_file=None, inputs=None, outputs=None, opset=None, channel_first=None): """ A command line interface for Keras/Tensorflow model to ONNX converter. :param input_file: the original model file path, could be a folder name of TF saved model :param output_file: the converted ONNX model file path (optional) :param inputs: The model graph input node list (tensorflow model only) :param outputs: The model graph output node list (tensorflow model only) :param opset: the target opset for the ONNX model. :param channel_first: the input name needs to be transposed as NCHW :return: """ if not os.path.exists(input_file): print("File or directory name '{}' is invalid!".format(input_file)) return file_ext = os.path.splitext(input_file) if output_file is None: output_file = file_ext[0] + '.onnx' v_channel_first = [] if channel_first is None else channel_first.split(sep=',') v_inputs = [] if inputs is None else inputs.split(sep=',') v_outputs = [] if outputs is None else outputs.split(sep=',') if file_ext[-1] == '.h5': kml = tf.keras.models.load_model(input_file) oxml = convert_keras(kml, kml.model, '', opset, channel_first) else: from tf2onnx import loader if os.path.isdir(input_file): chpt_file = [f_ for f_ in os.listdir(input_file) if f_.endswith('.meta')] if chpt_file: meta_file = os.path.join(input_file, chpt_file[0]) graph_def, v_inputs, v_outputs = loader.from_checkpoint(meta_file, v_inputs, v_outputs) else: saved_pbfile = [f_ for f_ in os.listdir(input_file) if f_.endswith('.pb')] assert len(saved_pbfile) == 1 saved_pbfile = os.path.join(input_file, saved_pbfile[0]) graph_def, v_inputs, v_outputs = loader.from_saved_model(saved_pbfile, v_inputs, v_outputs) else: with tf.gfile.GFile(input_file, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) oxml = convert_tensorflow(graph_def, os.path.basename(input_file), v_inputs, v_outputs, '', opset, v_channel_first) onnx.save_model(oxml, output_file)
def execute(self): if self._tf_format == 'TF_PB': graph_def = graph_pb2.GraphDef() with tf.gfile.GFile(self._tf_file, 'rb') as f: graph_def.ParseFromString(f.read()) inputs, outputs = _find_out_terminal_node(graph_def, postfix=True) else: if self._outputs is None: raise ImportError("Missing '--outputs' parameter.") if self._inputs is None: raise ImportError("Missing '--inputs' parameter.") inputs = [i + ":0" for i in self._inputs.split(",")] outputs = [i + ":0" for i in self._outputs.split(",")] if self._tf_format == 'TF_CKPT_V1': graph_def = self.load_checkpoint_v1() elif self._tf_format == 'TF_CKPT_V2': graph_def, inputs, outputs = loader.from_checkpoint( self._tf_file, inputs, outputs) onnx_model = self.convert_to_onnx(graph_def, inputs, outputs) onnx_importer = OnnxImporter() onnx_importer.import_from_onnx_model(onnx_model) return onnx_importer.execute()
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)
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.""" logger.info("===================================") logger.info("Running %s", name) self.perf = perf # get the model if self.url: _, dir_name = self.download_file() model_path = os.path.join(dir_name, self.local) else: model_path = self.local dir_name = os.path.dirname(self.local) logger.info("Downloaded to %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 = loader.from_checkpoint(model_path, input_names, outputs) elif self.model_type in ["saved_model"]: graph_def, input_names, outputs = loader.from_saved_model(model_path, input_names, outputs) else: graph_def, input_names, outputs = loader.from_graphdef(model_path, input_names, outputs) # create the input data inputs = {} for k, v in self.input_names.items(): if k not in input_names: continue if isinstance(v, six.text_type) and v.startswith("np."): inputs[k] = eval(v) # pylint: disable=eval-used else: inputs[k] = self.make_input(v) if self.more_inputs: for k, v in self.more_inputs.items(): inputs[k] = v graph_def = tf2onnx.tfonnx.tf_optimize(inputs.keys(), self.output_names, graph_def, fold_const) if utils.is_debug_mode(): utils.save_protobuf(os.path.join(TEMP_DIR, name + "_after_tf_optimize.pb"), graph_def) shape_override = {} g = tf.import_graph_def(graph_def, name='') with tf.Session(config=tf.ConfigProto(allow_soft_placement=True), graph=g) as sess: # fix inputs if needed for k in inputs.keys(): # pylint: disable=consider-iterating-dictionary t = sess.graph.get_tensor_by_name(k) dtype = tf.as_dtype(t.dtype).name v = inputs[k] if dtype != v.dtype: logger.warning("input dtype doesn't match tensorflow's") inputs[k] = np.array(v, dtype=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()) model_proto = onnx_graph.make_model("converted from tf2onnx") model_proto = optimizer.optimize_graph(onnx_graph).make_model("optimized") logger.info("To_ONNX, OK") if utils.is_debug_mode(): onnx_graph.dump_graph() if onnx_file: self.create_onnx_file(name, model_proto, inputs, onnx_file) except Exception: logger.error("To_ONNX FAIL", exc_info=1) try: onnx_results = None if backend == "caffe2": onnx_results = self.run_caffe2(name, model_proto, inputs) elif backend == "onnxmsrtnext": onnx_results = self.run_onnxmsrtnext(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 run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=None, extra_opset=None, perf=None, fold_const=None): """Run complete test against backend.""" print(name) self.perf = perf # get the model if self.url: _, dir_name = self.download_file() model_path = os.path.join(dir_name, self.local) else: model_path = self.local dir_name = os.path.dirname(self.local) print("\tdownloaded", model_path) inputs = list(self.input_names.keys()) outputs = self.output_names if self.model_type in ["checkpoint"]: graph_def, inputs, outputs = loader.from_checkpoint( model_path, inputs, outputs) elif self.model_type in ["saved_model"]: graph_def, inputs, outputs = loader.from_saved_model( model_path, inputs, outputs) else: graph_def, inputs, outputs = loader.from_graphdef( model_path, inputs, outputs) # create the input data inputs = {} for k, v in self.input_names.items(): if isinstance(v, six.text_type) and v.startswith("np."): inputs[k] = eval(v) # pylint: disable=eval-used else: inputs[k] = self.make_input(v) if self.more_inputs: for k, v in self.more_inputs.items(): inputs[k] = v graph_def = tf2onnx.tfonnx.tf_optimize(inputs.keys(), self.output_names, graph_def, fold_const) shape_override = {} g = tf.import_graph_def(graph_def, name='') with tf.Session(config=tf.ConfigProto(allow_soft_placement=True), graph=g) as sess: # fix inputs if needed for k in inputs.keys(): # pylint: disable=consider-iterating-dictionary t = sess.graph.get_tensor_by_name(k) dtype = tf.as_dtype(t.dtype).name v = inputs[k] if dtype != v.dtype: log.warning("input dtype doesn't match tensorflow's") inputs[k] = np.array(v, dtype=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: print("\ttensorflow", "SKIPPED") else: tf_results = self.run_tensorflow(sess, inputs) print("\ttensorflow", "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()) model_proto = onnx_graph.make_model("converted from tf2onnx") new_model_proto = optimizer.optimize_graph( onnx_graph, debug=debug).make_model("optimized") if new_model_proto: model_proto = new_model_proto else: print( "\tNON-CRITICAL, optimizers are not applied successfully" ) print("\tto_onnx", "OK") if debug: onnx_graph.dump_graph() if onnx_file: self.create_onnx_file(name, model_proto, inputs, onnx_file) except Exception as ex: tb = traceback.format_exc() print("\tto_onnx", "FAIL", ex, tb) try: onnx_results = None if backend == "caffe2": onnx_results = self.run_caffe2(name, model_proto, inputs) elif backend == "onnxmsrtnext": onnx_results = self.run_onnxmsrtnext(name, model_proto, inputs) elif backend == "onnxruntime": onnx_results = self.run_onnxruntime(name, model_proto, inputs) else: raise ValueError("unknown backend") print("\trun_onnx OK") try: if self.skip_tensorflow: print("\tResults: 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) print("\tResults: OK") return True except Exception as ex: print("\tResults: ", ex) except Exception as ex: print("\trun_onnx", "FAIL", ex) return False
def main(): args = get_args() # 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 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 # 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, 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=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 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 = loader.from_checkpoint(model_path, input_names, outputs) elif self.model_type in ["saved_model"]: graph_def, input_names, outputs = loader.from_saved_model(model_path, input_names, outputs) else: graph_def, input_names, outputs = loader.from_graphdef(model_path, input_names, outputs) # remove unused input names input_names = list(set(input_names).intersection(self.input_names.keys())) graph_def = tf2onnx.tfonnx.tf_optimize(input_names, self.output_names, graph_def, fold_const) if utils.is_debug_mode(): utils.save_protobuf(os.path.join(TEMP_DIR, name + "_after_tf_optimize.pb"), graph_def) inputs = {} shape_override = {} g = tf.import_graph_def(graph_def, name='') with tf.Session(config=tf.ConfigProto(allow_soft_placement=True), 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