def from_graph_def(graph_def, name=None, input_names=None, output_names=None, opset=None, custom_ops=None, custom_op_handlers=None, custom_rewriter=None, inputs_as_nchw=None, extra_opset=None, shape_override=None, target=None, large_model=False, tensors_to_rename=None, output_path=None): """Returns a ONNX model_proto for a tensorflow graphdef. Args: graph_def: the graphdef we want to convert input_names: list of input names output_names: list of output names name: A name for the graph opset: the opset to be used for the ONNX model, default is the latest custom_ops: if a model contains ops not recognized by onnx runtime, you can tag these ops with a custom op domain so that the runtime can still open the model. Type is a dictionary `{op name: domain}`. target: list of workarounds applied to help certain platforms custom_op_handlers: dictionary of custom ops handlers custom_rewriter: list of custom graph rewriters extra_opset: list of extra opset's, for example the opset's used by custom ops shape_override: dict with inputs that override the shapes given by tensorflow inputs_as_nchw: transpose inputs in list from nchw to nhwc large_model: use the ONNX external tensor storage format output_path: save model to output_path Returns: An ONNX model_proto and an external_tensor_storage dict. """ if not input_names: raise ValueError("input_names needs to be provided") if not output_names: raise ValueError("output_names needs to be provided") if not name: name = "unknown" initialized_tables = None with tf.device("/cpu:0"): with tf.Graph().as_default() as tf_graph: with tf_loader.tf_session(graph=tf_graph) as sess: tf.import_graph_def(graph_def, name='') frozen_graph = tf_loader.freeze_session(sess, input_names=input_names, output_names=output_names) input_names = tf_loader.inputs_without_resource(sess, input_names) frozen_graph = tf_loader.tf_optimize(input_names, output_names, graph_def) model_proto, external_tensor_storage = _convert_common( frozen_graph, name=name, continue_on_error=True, target=target, opset=opset, custom_op_handlers=custom_ops, extra_opset=extra_opset, shape_override=shape_override, input_names=input_names, output_names=output_names, inputs_as_nchw=inputs_as_nchw, large_model=large_model, tensors_to_rename=tensors_to_rename, initialized_tables=initialized_tables, output_path=output_path) return model_proto, external_tensor_storage
def _make_onnx_model(): with tf.compat.v1.Session(graph=tf.Graph()) as session: input_x = tf.compat.v1.placeholder(dtype=tf.float32, shape=[None, 2, 3, 4], name='x') input_y = tf.compat.v1.placeholder(dtype=tf.float32, shape=[None, 2, 3, 4], name='y') weight = tf.Variable(initial_value=4.2, dtype=tf.float32) tf.multiply(input_x + input_y, weight, name='z') session.run(weight.initializer) frozen_graph_def = tf_loader.freeze_session(session, input_names=['x:0', 'y:0'], output_names=['z:0']) with tf.compat.v1.Session(graph=tf.Graph()) as session: tf.import_graph_def(frozen_graph_def, name='') onnx_model = tfonnx.process_tf_graph( tf_graph=session.graph, input_names=['x:0', 'y:0'], output_names=['z:0']).make_model(graph_doc='Test onnx model') return OnnxModel(model_proto=onnx_model, input_data_formats=[ DataFormat.CHANNELS_LAST, DataFormat.CHANNELS_FIRST ])
def export_model(self, export_dir: str, onnx: Optional[int] = None) -> None: """Export tensorflow graph to export_dir for serving.""" if onnx: try: import tf2onnx except ImportError as e: raise RuntimeError( "Converting a TensorFlow model to ONNX requires " "`tf2onnx` to be installed. Install with " "`pip install tf2onnx`.") from e with self.get_session().graph.as_default(): signature_def_map = self._build_signature_def() sd = signature_def_map[ tf1.saved_model.signature_constants. DEFAULT_SERVING_SIGNATURE_DEF_KEY # noqa: E501 ] inputs = [v.name for k, v in sd.inputs.items()] outputs = [v.name for k, v in sd.outputs.items()] from tf2onnx import tf_loader frozen_graph_def = tf_loader.freeze_session( self._sess, input_names=inputs, output_names=outputs) with tf1.Session(graph=tf.Graph()) as session: tf.import_graph_def(frozen_graph_def, name="") g = tf2onnx.tfonnx.process_tf_graph( session.graph, input_names=inputs, output_names=outputs, inputs_as_nchw=inputs, ) model_proto = g.make_model("onnx_model") tf2onnx.utils.save_onnx_model(export_dir, "saved_model", feed_dict={}, model_proto=model_proto) else: with self.get_session().graph.as_default(): signature_def_map = self._build_signature_def() builder = tf1.saved_model.builder.SavedModelBuilder(export_dir) builder.add_meta_graph_and_variables( self.get_session(), [tf1.saved_model.tag_constants.SERVING], signature_def_map=signature_def_map, saver=tf1.summary.FileWriter(export_dir).add_graph( graph=self.get_session().graph), ) builder.save()
def load(self, model_path: Union[str, Path], **_) -> Model: if isinstance(model_path, Path): model_path = model_path.as_posix() get_model = load_from_file(model_path, "model", GET_MODEL_FN_NAME) get_serving_input_receiver_fn = load_from_file(model_path, "model", GET_SERVING_INPUT_RECEIVER_FN) if get_model is None: raise RuntimeError(f"Could not find {GET_MODEL_FN_NAME} in {model_path}") if get_serving_input_receiver_fn is None: raise RuntimeError(f"Could not find {GET_SERVING_INPUT_RECEIVER_FN} in {model_path}") model_args = filter_fn_args(self._model_args, fn=get_model) serving_input_receiver_args = filter_fn_args(self._model_args, fn=get_serving_input_receiver_fn) session_config = create_session_config(allow_growth=True) tf.compat.v1.reset_default_graph() with tf.compat.v1.Session(config=session_config) as sess: estimator = get_model(**model_args) serving_input_receiver_fn = get_serving_input_receiver_fn(**serving_input_receiver_args) input_receiver = serving_input_receiver_fn() estimator_spec = estimator.model_fn( features=input_receiver.features, labels=None, mode=tf.estimator.ModeKeys.PREDICT, config=estimator.config, ) input_tensors_dict = input_receiver.receiver_tensors output_tensors_dict = estimator_spec.predictions inputs_dict = {k: tensor2tensor_spec(tensor) for k, tensor in input_tensors_dict.items()} outputs_dict = {k: tensor2tensor_spec(tensor) for k, tensor in output_tensors_dict.items()} input_tensor_names = [t.name for t in inputs_dict.values()] output_tensor_names = [t.name for t in outputs_dict.values()] graph_saver = estimator_spec.scaffold.saver or tf.compat.v1.train.Saver(sharded=True) graph_saver.restore(sess, estimator.latest_checkpoint()) input_tensor_names = inputs_without_resource(sess, input_tensor_names) frozen_graph = freeze_session(sess, input_names=input_tensor_names, output_names=output_tensor_names) input_tensor_names = remove_redundant_inputs(frozen_graph, input_tensor_names) tf.compat.v1.reset_default_graph() with tf.compat.v1.Session(config=estimator.config.session_config): frozen_graph = tf_optimize(input_tensor_names, output_tensor_names, frozen_graph) tf.compat.v1.reset_default_graph() precision = _infer_model_precision(frozen_graph, inputs_dict, outputs_dict) return Model(frozen_graph, precision, inputs_dict, outputs_dict)
def _from_saved_model_v1(sess, model_path, tag, signatures): """ Load tensorflow graph from saved_model. NOTICE: Modified version from tf2onnx project """ wrn_no_tag = "'--tag' not specified for saved_model. Using --tag serve" wrn_empty_tag = "'--tag' value is empty string. Using tag =[[]]" if tag is None: tag = [tf.saved_model.SERVING] LOGGER.warning(wrn_no_tag) if tag == "": tag = [[]] LOGGER.warning(wrn_empty_tag) if not isinstance(tag, list): tag = [tag] imported = tf.compat.v1.saved_model.loader.load(sess, tag, model_path) for k in imported.signature_def.keys(): if k.startswith("_"): # consider signatures starting with '_' private continue signatures.append(k) try: from tensorflow.contrib.saved_model.python.saved_model import ( # pytype: disable=import-error signature_def_utils, ) # pylint: disable=unnecessary-lambda get_signature_def = lambda meta_graph_def, k: signature_def_utils.get_signature_def_by_key( meta_graph_def, k) except ImportError: # TF1.12 changed the api get_signature_def = lambda meta_graph_def, k: meta_graph_def.signature_def[ k] inputs = {} outputs = {} for k in signatures: inputs_tensor_info = get_signature_def(imported, k).inputs for name, input_tensor in inputs_tensor_info.items(): inputs[name] = input_tensor.name outputs_tensor_info = get_signature_def(imported, k).outputs for name, output_tensor in outputs_tensor_info.items(): outputs[name] = output_tensor.name frozen_graph = freeze_session(sess, input_names=list(inputs.values()), output_names=list(outputs.values())) return frozen_graph, inputs, outputs
def _from_keras_tf1(model, input_signature=None, opset=None, custom_ops=None, custom_op_handlers=None, custom_rewriter=None, inputs_as_nchw=None, extra_opset=None, shape_override=None, target=None, large_model=False, output_path=None): """from_keras for tf 1.15""" input_names = [t.name for t in model.inputs] output_names = [t.name for t in model.outputs] old_out_names = _rename_duplicate_keras_model_names(model) tensors_to_rename = dict(zip(input_names, model.input_names)) tensors_to_rename.update(zip(output_names, model.output_names)) if old_out_names is not None: model.output_names = old_out_names if _is_legacy_keras_model(model): import keras # pylint: disable=import-outside-toplevel sess = keras.backend.get_session() else: sess = tf.keras.backend.get_session(model.outputs) with tf.device("/cpu:0"): frozen_graph, initialized_tables = tf_loader.freeze_session(sess, input_names, output_names, get_tables=True) with tf.Graph().as_default(): tf.import_graph_def(frozen_graph, name="") frozen_graph = tf_loader.tf_optimize(input_names, output_names, frozen_graph, False) model_proto, external_tensor_storage = _convert_common( frozen_graph, name=model.name, continue_on_error=True, target=target, opset=opset, custom_ops=custom_ops, custom_op_handlers=custom_op_handlers, custom_rewriter=custom_rewriter, extra_opset=extra_opset, shape_override=shape_override, input_names=input_names, output_names=output_names, inputs_as_nchw=inputs_as_nchw, large_model=large_model, tensors_to_rename=tensors_to_rename, initialized_tables=initialized_tables, output_path=output_path) return model_proto, external_tensor_storage
def _save_onnx_model(model_file): with tf.compat.v1.Session(graph=tf.Graph()) as session: input_x = tf.compat.v1.placeholder(dtype=tf.float32, shape=[None, 2, 3, 4], name='x') input_y = tf.compat.v1.placeholder(dtype=tf.float32, shape=[None, 2, 3, 4], name='y') weight = tf.Variable(initial_value=4.2, dtype=tf.float32) tf.multiply(input_x + input_y, weight, name='z') session.run(weight.initializer) frozen_graph_def = tf_loader.freeze_session(session, input_names=['x:0', 'y:0'], output_names=['z:0']) with tf.compat.v1.Session(graph=tf.Graph()) as session, open( model_file.name, mode='wb') as graph_file: tf.import_graph_def(frozen_graph_def, name='') onnx_model = tfonnx.process_tf_graph( tf_graph=session.graph, input_names=['x:0', 'y:0'], output_names=['z:0']).make_model(graph_doc='Test onnx model') graph_file.write(onnx_model.SerializeToString())
def run_test_case(self, func, feed_dict, input_names_with_port, output_names_with_port, rtol=1e-07, atol=1e-5, convert_var_to_const=True, constant_fold=True, check_value=True, check_shape=True, check_dtype=True, process_args=None, onnx_feed_dict=None, graph_validator=None, as_session=False, large_model=False): # optional - passed to process_tf_graph if process_args is None: process_args = {} # optional - pass distinct feed_dict to onnx runtime if onnx_feed_dict is None: onnx_feed_dict = feed_dict input_names_with_port = list(feed_dict) tf_reset_default_graph() graph_def = None np.random.seed(1) # Make it reproducible. clean_feed_dict = {utils.node_name(k): v for k, v in feed_dict.items()} if is_tf2() and not as_session: # # use eager to execute the tensorflow func # # numpy doesn't work for all ops, make it tf.Tensor() input_tensors = [ tf.TensorSpec(shape=v.shape, dtype=tf.as_dtype(v.dtype), name=utils.node_name(k)) for k, v in feed_dict.items() ] input_list = [ tf.convert_to_tensor(v, dtype=tf.as_dtype(v.dtype), name=utils.node_name(k)) for k, v in feed_dict.items() ] tf.random.set_seed(1) expected = func(*input_list) if isinstance(expected, (list, tuple)): # list or tuple expected = [x.numpy() for x in expected] else: # single result expected = [expected.numpy()] # now make the eager functions a graph concrete_func = tf.function(func, input_signature=tuple(input_tensors)) concrete_func = concrete_func.get_concrete_function() graph_def = from_function(concrete_func, input_names=list(feed_dict.keys()), output_names=output_names_with_port, large_model=large_model) else: # # use graph to execute the tensorflow func # with tf_session() as sess: tf_set_random_seed(1) input_list = [] for k, v in clean_feed_dict.items(): input_list.append( tf_placeholder(name=k, shape=v.shape, dtype=tf.as_dtype(v.dtype))) func(*input_list) variables_lib.global_variables_initializer().run() tf_tables_initializer().run() output_dict = [] for out_name in output_names_with_port: output_dict.append(sess.graph.get_tensor_by_name(out_name)) expected = sess.run(output_dict, feed_dict=feed_dict) graph_def = freeze_session(sess, input_names=list(feed_dict.keys()), output_names=output_names_with_port) tf_reset_default_graph() with tf_session() as sess: tf.import_graph_def(graph_def, name='') graph_def = tf_optimize(list(feed_dict.keys()), output_names_with_port, graph_def, fold_constant=constant_fold) tf_reset_default_graph() with tf_session() as sess: const_node_values = None if large_model: const_node_values = compress_graph_def(graph_def) tf.import_graph_def(graph_def, name='') if self.config.is_debug_mode: model_path = os.path.join( self.test_data_directory, self._testMethodName + "_after_tf_optimize.pb") utils.save_protobuf(model_path, graph_def) self.logger.debug("created file %s", model_path) g = process_tf_graph(sess.graph, opset=self.config.opset, input_names=list(feed_dict.keys()), output_names=output_names_with_port, target=self.config.target, const_node_values=const_node_values, **process_args) g = optimizer.optimize_graph(g) actual = self.run_backend(g, output_names_with_port, onnx_feed_dict, large_model) for expected_val, actual_val in zip(expected, actual): if check_value: self.assertAllClose(expected_val, actual_val, rtol=rtol, atol=atol) if check_dtype: self.assertEqual(expected_val.dtype, actual_val.dtype) # why need shape checke: issue when compare [] with scalar # https://github.com/numpy/numpy/issues/11071 if check_shape: self.assertEqual(expected_val.shape, actual_val.shape) if graph_validator: self.assertTrue(graph_validator(g)) return g
def freeze_and_run_tf(self, func, feed_dict, outputs, as_session, premade_placeholders, large_model, constant_fold): np.random.seed(1) # Make it reproducible. clean_feed_dict = {utils.node_name(k): v for k, v in feed_dict.items()} if is_tf2() and not as_session: # # use eager to execute the tensorflow func # # numpy doesn't work for all ops, make it tf.Tensor() input_tensors = [ tf.TensorSpec(shape=v.shape, dtype=tf.as_dtype(v.dtype), name=utils.node_name(k)) for k, v in feed_dict.items() ] input_list = [ tf.convert_to_tensor(v, dtype=tf.as_dtype(v.dtype), name=utils.node_name(k)) for k, v in feed_dict.items() ] tf.random.set_seed(1) result = func(*input_list) if isinstance(result, (list, tuple)): # list or tuple result = [x.numpy() for x in result] else: # single result result = [result.numpy()] # now make the eager functions a graph concrete_func = tf.function(func, input_signature=tuple(input_tensors)) concrete_func = concrete_func.get_concrete_function() graph_def = from_function(concrete_func, input_names=list(feed_dict.keys()), output_names=outputs, large_model=large_model) initialized_tables = None else: # # use graph to execute the tensorflow func # with tf_session() as sess: tf_set_random_seed(1) input_list = [] if not premade_placeholders: for k, v in clean_feed_dict.items(): input_list.append( tf_placeholder(name=k, shape=v.shape, dtype=tf.as_dtype(v.dtype))) func(*input_list) variables_lib.global_variables_initializer().run() tf_tables_initializer().run() output_dict = [] for out_name in outputs: output_dict.append(sess.graph.get_tensor_by_name(out_name)) result = sess.run(output_dict, feed_dict=feed_dict) graph_def = freeze_session(sess, input_names=list(feed_dict.keys()), output_names=outputs) table_names, key_dtypes, value_dtypes = get_hash_table_info( graph_def) initialized_tables = {} for n, k_dtype, val_dtype in zip(table_names, key_dtypes, value_dtypes): h = lookup_ops.hash_table_v2(k_dtype, val_dtype, shared_name=n) k, v = lookup_ops.lookup_table_export_v2( h, k_dtype, val_dtype) initialized_tables[n] = (sess.run(k), sess.run(v)) tf_reset_default_graph() with tf_session() as sess: tf.import_graph_def(graph_def, name='') graph_def = tf_optimize(list(feed_dict.keys()), outputs, graph_def, fold_constant=constant_fold) model_path = os.path.join( self.test_data_directory, self._testMethodName + "_after_tf_optimize.pb") utils.save_protobuf(model_path, graph_def) self.logger.debug("created file %s", model_path) return result, graph_def, initialized_tables
def freeze_and_run_tf(self, func, feed_dict, outputs, as_session, premade_placeholders, large_model): np.random.seed(1) # Make it reproducible. clean_feed_dict = {utils.node_name(k): v for k, v in feed_dict.items()} if is_tf2() and not as_session: # # use eager to execute the tensorflow func # # numpy doesn't work for all ops, make it tf.Tensor() input_tensors = [ tf.TensorSpec(shape=v.shape, dtype=tf.as_dtype(v.dtype), name=utils.node_name(k)) for k, v in feed_dict.items() ] input_list = [ tf.convert_to_tensor(v, dtype=tf.as_dtype(v.dtype), name=utils.node_name(k)) for k, v in feed_dict.items() ] tf.random.set_seed(1) result = func(*input_list) if isinstance(result, (list, tuple)): # list or tuple result = [x.numpy() for x in result] else: # single result result = [result.numpy()] # now make the eager functions a graph concrete_func = tf.function(func, input_signature=tuple(input_tensors)) concrete_func = concrete_func.get_concrete_function() graph_def = from_function(concrete_func, input_names=list(feed_dict.keys()), output_names=outputs, large_model=large_model) initialized_tables = None else: # # use graph to execute the tensorflow func # with tf_session() as sess: tf_set_random_seed(1) input_list = [] if not premade_placeholders: for k, v in clean_feed_dict.items(): input_list.append( tf_placeholder(name=k, shape=v.shape, dtype=tf.as_dtype(v.dtype))) func(*input_list) variables_lib.global_variables_initializer().run() tf_tables_initializer().run() output_dict = [] for out_name in outputs: output_dict.append(sess.graph.get_tensor_by_name(out_name)) result = sess.run(output_dict, feed_dict=feed_dict) graph_def = freeze_session(sess, input_names=list(feed_dict.keys()), output_names=outputs) table_info = get_hash_table_info(graph_def) initialized_tables = {} for info in table_info: if info.shared_name is None: continue h = lookup_ops.hash_table_v2(info.key_dtype, info.val_dtype, shared_name=info.shared_name) k, v = lookup_ops.lookup_table_export_v2( h, info.key_dtype, info.val_dtype) initialized_tables[info.shared_name] = (sess.run(k), sess.run(v)) tf_reset_default_graph() with tf_session() as sess: tf.import_graph_def(graph_def, name='') graph_def = tf_optimize(list(feed_dict.keys()), outputs, graph_def) return result, graph_def, initialized_tables