def compile_source(source: TensorFlowFrozenGraphModel) -> OnnxModel: with tf.Graph().as_default() as graph: tf.import_graph_def(graph_def=source.graph_def, name='') onnx_graph = tfonnx.process_tf_graph( tf_graph=graph, opset=min(10, defs.onnx_opset_version()), input_names=[source_input.name for source_input in source.inputs], output_names=list(source.outputs)) onnx_graph = optimizer.optimize_graph(graph=onnx_graph) model_proto = onnx_graph.make_model( 'ONNX model generated from TensorFlow frozen graph model.') input_name_to_index = { source_input.name: i for i, source_input in enumerate(source.inputs) } sorted_graph_inputs = sorted( model_proto.graph.input, key=lambda target_input: input_name_to_index[target_input.name]) del model_proto.graph.input[:] model_proto.graph.input.extend(sorted_graph_inputs) return OnnxModel(model_proto=model_proto, input_data_formats=[ source_input.data_format for source_input in source.inputs ])
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 compile_source(source: TfModelFile, config: Config) -> OnnxModel: graph_def, inputs, outputs = tf_loader.from_checkpoint( source.model_path + '.meta', config.input_names, config.output_names) with tf.Graph().as_default() as graph: tf.import_graph_def(graph_def=graph_def, name='') onnx_graph = tfonnx.process_tf_graph(tf_graph=graph, opset=min(10, defs.onnx_opset_version()), input_names=inputs, output_names=outputs) onnx_graph = optimizer.optimize_graph(graph=onnx_graph) model_proto = onnx_graph.make_model( 'ONNX model generated from TensorFlow model.') input_name_to_index = { source_input: i for i, source_input in enumerate(inputs) } sorted_graph_inputs = sorted( model_proto.graph.input, key=lambda target_input: input_name_to_index[target_input.name]) del model_proto.graph.input[:] model_proto.graph.input.extend(sorted_graph_inputs) return OnnxModel(model_proto=model_proto, input_data_formats=config.data_formats)
def tfgraph2onnx(graph_def, inputnames2tensornames, outputnames2tensornames, *, onnx_opset, onnx_optimized=True): with tf.Graph().as_default() as tf_graph: tf.import_graph_def(graph_def, name="") session_config = create_session_config(allow_growth=True) with tf.compat.v1.Session(graph=tf_graph, config=session_config): input_tensor_names = list(inputnames2tensornames.values()) output_tensor_names = list(outputnames2tensornames.values()) onnx_graph = tfonnx.process_tf_graph( tf_graph, input_names=input_tensor_names, output_names=output_tensor_names, opset=onnx_opset, ) if onnx_optimized: onnx_graph = optimizer.optimize_graph(onnx_graph) graph_doc: str = "triton export" onnx_model = onnx_graph.make_model(graph_doc) # to match tensorflow savedmodel signature _replace_io_names(onnx_model.graph, "inputs", inputnames2tensornames) _replace_io_names(onnx_model.graph, "outputs", outputnames2tensornames) onnx.checker.check_model(onnx_model) onnx.helper.strip_doc_string(onnx_model) onnx_model = onnx.shape_inference.infer_shapes(onnx_model) return onnx_model
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_common(frozen_graph, name="unknown", large_model=False, output_path=None, output_frozen_graph=None, **kwargs): """Common processing for conversion.""" model_proto = None external_tensor_storage = None const_node_values = None with tf.Graph().as_default() as tf_graph: if large_model: const_node_values = compress_graph_def(frozen_graph) external_tensor_storage = ExternalTensorStorage() if output_frozen_graph: utils.save_protobuf(output_frozen_graph, frozen_graph) if not kwargs.get("tflite_path"): tf.import_graph_def(frozen_graph, name='') g = process_tf_graph(tf_graph, const_node_values=const_node_values, **kwargs) onnx_graph = optimizer.optimize_graph(g, catch_errors=not large_model) model_proto = onnx_graph.make_model( "converted from {}".format(name), external_tensor_storage=external_tensor_storage) if output_path: if large_model: utils.save_onnx_zip(output_path, model_proto, external_tensor_storage) else: utils.save_protobuf(output_path, model_proto) return model_proto, external_tensor_storage
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 convert(input_ops_dict, output_ops, input_model, output_model): '''Convert keras h5 to tensorflow pb Args: input_ops_dict: input ops dict including names and shapes output_ops: output op names input_model: input keras h5 model name output_model: output pb model name ''' onnx_name = ".tmp.onnx" pb_name = ".tmp.pb" # keras --> onnx --> pb --> onnx --> pb # keras --> onnx model = tf.keras.models.load_model(input_model) onnx_model = keras2onnx.convert_keras(model, model.name, target_opset=8) keras2onnx.save_model(onnx_model, onnx_name) # onnx --> tf onnx_model = onnx.load(onnx_name) tf_rep = prepare(onnx_model, input_shape_dict=input_ops_dict) tf_rep.export_graph(pb_name) # tf --> onnx (fold constants) inputs = input_ops_dict.keys() inputs = [i + ":0" for i in inputs] outputs = output_ops graph_def, inputs, outputs = tf_loader.from_graphdef( pb_name, inputs, outputs) with tf.Graph().as_default() as tf_graph: tf.import_graph_def(graph_def, name="") g = tf2onnx.tfonnx.process_tf_graph(tf_graph, opset=8, input_names=inputs, output_names=outputs) onnx_graph = optimizer.optimize_graph(g) model_proto = onnx_graph.make_model("converted from %s" % pb_name) utils.save_protobuf(onnx_name, model_proto) # onnx --> tf onnx_model = onnx.load(onnx_name) tf_rep = prepare(onnx_model, input_shape_dict=input_ops_dict) tf_rep.export_graph(output_model) # remove tmp files if os.path.exists(onnx_name): os.remove(onnx_name) if os.path.exists(pb_name): os.remove(pb_name)
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 convert_frozen_pb_to_onnx(frozen_pb_or_graph_def, opset=9, tf_graph_optimization=True, input_shape=None, name=None): try: from tf2onnx.tfonnx import process_tf_graph, tf_optimize from tf2onnx import constants, loader, logging, utils, optimizer except Exception as e: logger.error('import tf2onnx error, "pip install tf2onnx"') exit(0) graph_def = frozen_pb_or_graph_def if isinstance(frozen_pb_or_graph_def, str): model_path = frozen_pb_or_graph_def output_dir = model_path.replace('.pb', '.onnx') else: model_path = 'graphdef_buffer' assert name, 'name should be give to export an .onnx when converting from a graphdef buffer' output_dir = '{}.onnx'.format(name) inputs, outputs = _auto_inputs_outputs_detect(graph_def) shape_override = {} if input_shape: assert isinstance(input_shape, list), 'input_shape item need to be list, each for dims of a input tensor' for idx, item in enumerate(input_shape): shape_override[inputs[idx]] = item # graph optimizatin with tf_graph_transform if tf_graph_optimization: graph_def = graph_optimization(graph_def) 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='', opset=opset, custom_op_handlers={}, extra_opset=[], shape_override=shape_override, input_names=inputs, output_names=outputs, inputs_as_nchw=None) # graph optimization with onnx optimizer 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) outputs = model_path.replace('.pb', '.onnx') utils.save_protobuf(output_dir, model_proto) logger.info("ONNX model is saved at %s", output_dir)
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 run_tflite_test(self, tflite_model, feed_dict, rtol=1e-07, atol=1e-5): 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) tf_lite_output_data, output_names = self.run_tflite(tflite_path, feed_dict) g = process_tf_graph(None, opset=self.config.opset, input_names=list(feed_dict.keys()), output_names=output_names, target=self.config.target, tflite_path=tflite_path) g = optimizer.optimize_graph(g) onnx_from_tfl_output = self.run_backend(g, output_names, feed_dict, postfix="_from_tflite") self.assert_results_equal(tf_lite_output_data, onnx_from_tfl_output, rtol, atol)
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 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 _convert_common(frozen_graph, name="unknown", large_model=False, output_path=None, output_frozen_graph=None, custom_ops=None, custom_op_handlers=None, **kwargs): """Common processing for conversion.""" model_proto = None external_tensor_storage = None const_node_values = None if custom_ops is not None: if custom_op_handlers is None: custom_op_handlers = {} custom_op_handlers.update( {op: (make_default_custom_op_handler(domain), []) for op, domain in custom_ops.items()}) with tf.Graph().as_default() as tf_graph: if large_model: const_node_values = compress_graph_def(frozen_graph) external_tensor_storage = ExternalTensorStorage() if output_frozen_graph: utils.save_protobuf(output_frozen_graph, frozen_graph) if not kwargs.get("tflite_path") and not kwargs.get("tfjs_path"): tf.import_graph_def(frozen_graph, name='') g = process_tf_graph(tf_graph, const_node_values=const_node_values, custom_op_handlers=custom_op_handlers, **kwargs) if constants.ENV_TF2ONNX_CATCH_ERRORS in os.environ: catch_errors = constants.ENV_TF2ONNX_CATCH_ERRORS.upper() == "TRUE" else: catch_errors = not large_model onnx_graph = optimizer.optimize_graph(g, catch_errors) model_proto = onnx_graph.make_model("converted from {}".format(name), external_tensor_storage=external_tensor_storage) if output_path: if large_model: utils.save_onnx_zip(output_path, model_proto, external_tensor_storage) else: utils.save_protobuf(output_path, model_proto) return model_proto, external_tensor_storage
def main(): args = get_args() g, org_model_proto = load_graph(args.input, args.target) if g.is_target(constants.TARGET_CHANNELS_FIRST): g.reset_nodes(rewrite_channels_first(g, g.get_nodes())) if g.is_target(constants.TARGET_CHANNELS_LAST): g.reset_nodes(rewrite_channels_last(g, g.get_nodes())) g = optimizer.optimize_graph(g) onnx_graph = g.make_graph(org_model_proto.graph.doc_string + " (+tf2onnx/onnx-optimize)") kwargs = GraphUtil.get_onnx_model_properties(org_model_proto) model_proto = helper.make_model(onnx_graph, **kwargs) # write onnx graph if args.output: with open(args.output, "wb") as f: f.write(model_proto.SerializeToString())
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="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 optimize_graph(graph): return optimizer.optimize_graph(graph)
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 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(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 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) else: graph_def, input_names, outputs = tf_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
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: # 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, 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, 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_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 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 [] tflite_path = None 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) if any(opset.domain == constants.CONTRIB_OPS_DOMAIN for opset in extra_opset): try: import tensorflow_text # pylint: disable=import-outside-toplevel except ModuleNotFoundError: logger.warning( "tensorflow_text not installed. Model will fail to load if tensorflow_text ops are used." ) # 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.tflite: graph_def = None inputs = None outputs = None tflite_path = args.tflite model_path = tflite_path if args.verbose: logger.info("inputs: %s", inputs) logger.info("outputs: %s", outputs) tf_graph = None const_node_values = None if graph_def is not None: 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, tflite_path=tflite_path, dequantize=args.dequantize) 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_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 optimize_graph(graph, debug=False): return optimizer.optimize_graph(graph, debug)
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