def from_function(function, 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): """Returns a ONNX model_proto for a tf.function. Args: function: the tf.function we want to convert input_signature: a tf.TensorSpec or a numpy array defining the shape/dtype of the input 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 LooseVersion(tf.__version__) < "2.0": raise NotImplementedError("from_function requires tf-2.0 or newer") if not input_signature: raise ValueError("from_function requires input_signature") concrete_func = function.get_concrete_function(*input_signature) input_names = [input_tensor.name for input_tensor in concrete_func.inputs if input_tensor.dtype != tf.dtypes.resource] output_names = [output_tensor.name for output_tensor in concrete_func.outputs if output_tensor.dtype != tf.dtypes.resource] initialized_tables = None tensors_to_rename = tensor_names_from_structed(concrete_func, input_names, output_names) with tf.device("/cpu:0"): frozen_graph = tf_loader.from_function(concrete_func, input_names, output_names, large_model=large_model) model_proto, external_tensor_storage = _convert_common( frozen_graph, name=concrete_func.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 from_keras(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): """Returns a ONNX model_proto for a tf.keras model. Args: model: the tf.keras model we want to convert input_signature: a tf.TensorSpec or a numpy array defining the shape/dtype of the input opset: the opset to be used for the ONNX model, default is the latest 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 LooseVersion(tf.__version__) < "2.0": raise NotImplementedError("from_keras requires tf-1.15 or newer") from tensorflow.python.keras.saving import saving_utils as _saving_utils # pylint: disable=import-outside-toplevel # let tensorflow do the checking if model is a valid model function = _saving_utils.trace_model_call(model, input_signature) concrete_func = function.get_concrete_function(*input_signature) input_names = [input_tensor.name for input_tensor in concrete_func.inputs if input_tensor.dtype != tf.dtypes.resource] output_names = [output_tensor.name for output_tensor in concrete_func.outputs if output_tensor.dtype != tf.dtypes.resource] initialized_tables = None tensors_to_rename = tensor_names_from_structed(concrete_func, input_names, output_names) with tf.device("/cpu:0"): frozen_graph = tf_loader.from_function(concrete_func, input_names, output_names, large_model=large_model) model_proto, external_tensor_storage = _convert_common( frozen_graph, name=model.name, continue_on_error=True, target=None, 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 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) if get_model is None: raise RuntimeError(f"Could not find {GET_MODEL_FN_NAME} in {model_path}") model_args = filter_fn_args(self._model_args, fn=get_model) if self._allow_growth: physical_devices = tf.config.experimental.list_physical_devices("GPU") for device in physical_devices: tf.config.experimental.set_memory_growth(device, True) tf.keras.backend.clear_session() tf.keras.backend.set_learning_phase(False) eager_model, call_fn = get_model(**model_args) inputs_dict: Dict[str, TensorSpec] = { input_name: TensorSpec(t.name, t.dtype.name, tuple(t.shape.as_list())) for input_name, t in zip(eager_model.input_names, eager_model.inputs) } concrete_func = call_fn.get_concrete_function( *[tf.TensorSpec(shape=spec.shape, dtype=spec.dtype, name=name) for name, spec in inputs_dict.items()] ) input_tensors_names = [tensor.name for tensor in concrete_func.inputs if tensor.dtype != tf.dtypes.resource] output_tensors_names = [tensor.name for tensor in concrete_func.outputs] graph_def = from_function( concrete_func, input_tensors_names, output_tensors_names, large_model=self._large_model ) # tensor names changes after wrapping with call_fn, thus need to use those from concrete_func outputs_dict: Dict[str, TensorSpec] = { output_name: TensorSpec(output_tensor_name, t.dtype.name, tuple(t.shape.as_list())) for output_name, output_tensor_name, t in zip( eager_model.output_names, output_tensors_names, eager_model.outputs ) } precision = _infer_model_precision(graph_def, inputs_dict, outputs_dict) tf.keras.backend.clear_session() tf.keras.backend.set_learning_phase(False) def _add_suffix_as_quickfix_for_tf24_func_refactor(spec): if not spec.name.endswith(":0"): spec = spec._replace(name=spec.name + ":0") return spec inputs_dict = {name: _add_suffix_as_quickfix_for_tf24_func_refactor(spec) for name, spec in inputs_dict.items()} return Model(graph_def, precision, inputs_dict, outputs_dict)
def test_parse_tflite_graph(self): def func(a, b, c): alpha = tf.constant(1.1, dtype=tf.float32) beta = tf.constant(2.3, dtype=tf.float32) mul1 = tf.multiply(alpha, tf.matmul(a, b)) mul2 = tf.multiply(beta, c) x_ = mul1 + mul2 return tf.identity(x_, name="output") inp_shapes = [[2, 3], [3, 1], [2, 1]] inp_dtypes = [tf.float32, tf.float32, tf.float32] names = ['a', 'b', 'c'] names_with_port = ['a:0', 'b:0', 'c:0'] output_names = ['output'] output_names_with_port = ['output:0'] input_tensors = [tf.TensorSpec(shape=s, dtype=d, name=n) for s, d, n in zip(inp_shapes, inp_dtypes, names)] concrete_func = tf.function(func, input_signature=tuple(input_tensors)) concrete_func = concrete_func.get_concrete_function() graph_def = from_function(concrete_func, input_names=names_with_port, output_names=output_names_with_port) with tf_session() as sess: tf.import_graph_def(graph_def, name='') sess_inputs = [sess.graph.get_tensor_by_name(k) for k in names_with_port] sess_outputs = [sess.graph.get_tensor_by_name(n) for n in output_names_with_port] converter = tf.compat.v1.lite.TFLiteConverter.from_session(sess, sess_inputs, sess_outputs) tflite_model = converter.convert() tflite_path = os.path.join(self.test_data_directory, self._testMethodName + ".tflite") dir_name = os.path.dirname(tflite_path) tflite_model = converter.convert() os.makedirs(dir_name, exist_ok=True) with open(tflite_path, 'wb') as f: f.write(tflite_model) tflite_graphs, opcodes_map, model, tensor_shapes = read_tflite_model(tflite_path) self.assertEqual(1, len(tflite_graphs)) onnx_nodes, op_cnt, attr_cnt, output_shapes, dtypes, inputs, outputs, _ = \ parse_tflite_graph(tflite_graphs[0], opcodes_map, model, tensor_shapes_override=tensor_shapes) self.assertEqual(2, op_cnt['MUL']) self.assertEqual(1, op_cnt['ADD']) self.assertEqual(1, op_cnt['FULLY_CONNECTED']) self.assertEqual(1, attr_cnt['WeightsFormat']) self.assertEqual(names, inputs) self.assertEqual(output_names, outputs) for name, shape, dtype in zip(names, inp_shapes, inp_dtypes): self.assertEqual(shape, output_shapes[name]) self.assertEqual(dtype, dtypes[name]) self.assertTrue(len(onnx_nodes) >= 4)
def 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 from_keras(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): """Returns a ONNX model_proto for a tf.keras model. Args: model: the tf.keras model we want to convert input_signature: a tf.TensorSpec or a numpy array defining the shape/dtype of the input opset: the opset to be used for the ONNX model, default is the latest 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 LooseVersion(tf.__version__) < "2.0": raise NotImplementedError("from_keras requires tf-2.0 or newer") from tensorflow.python.keras.saving import saving_utils as _saving_utils # pylint: disable=import-outside-toplevel # let tensorflow do the checking if model is a valid model function = _saving_utils.trace_model_call(model, input_signature) try: concrete_func = function.get_concrete_function() except TypeError as e: # Legacy keras models don't accept the training arg tf provides so we hack around it if "got an unexpected keyword argument 'training'" not in str(e): raise e model_call = model.call def wrap_call(*args, training=False, **kwargs): return model_call(*args, **kwargs) model.call = wrap_call function = _saving_utils.trace_model_call(model, input_signature) concrete_func = function.get_concrete_function() # Put it back model.call = model_call # These inputs will be removed during freezing (includes resources, etc.) graph_captures = concrete_func.graph._captures # pylint: disable=protected-access captured_inputs = [ t_name.name for t_val, t_name in graph_captures.values() ] input_names = [ input_tensor.name for input_tensor in concrete_func.inputs if input_tensor.name not in captured_inputs ] output_names = [ output_tensor.name for output_tensor in concrete_func.outputs if output_tensor.dtype != tf.dtypes.resource ] initialized_tables = None tensors_to_rename = tensor_names_from_structed(concrete_func, input_names, output_names) reverse_lookup = {v: k for k, v in tensors_to_rename.items()} if model.output_names: # model.output_names is an optional field of Keras models indicating output order. It is None if unused. output_names = [reverse_lookup[out] for out in model.output_names] elif isinstance(concrete_func.structured_outputs, dict): # Other models specify output order using the key order of structured_outputs output_names = [ reverse_lookup[out] for out in concrete_func.structured_outputs.keys() ] with tf.device("/cpu:0"): frozen_graph = tf_loader.from_function(concrete_func, input_names, output_names, large_model=large_model) model_proto, external_tensor_storage = _convert_common( frozen_graph, name=model.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 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