def test_specify_input_signature(self): model = testing_utils.get_small_sequential_mlp(10, 3, None) inputs = array_ops.ones((8, 5)) with self.assertRaisesRegexp(ValueError, 'input shapes have not been set'): saving_utils.trace_model_call(model) fn = saving_utils.trace_model_call( model, [tensor_spec.TensorSpec(shape=[None, 5], dtype=dtypes.float32)]) signature_outputs = fn(inputs) expected_outputs = {model.output_names[0]: model(inputs)} self._assert_all_close(expected_outputs, signature_outputs)
def test_trace_model_outputs(self): input_dim = 5 if testing_utils.get_model_type() == 'functional' else None model = testing_utils.get_small_mlp(10, 3, input_dim) inputs = array_ops.ones((8, 5)) if input_dim is None: with self.assertRaisesRegexp(ValueError, 'input shapes have not been set'): saving_utils.trace_model_call(model) model._set_inputs(inputs) fn = saving_utils.trace_model_call(model) signature_outputs = fn(inputs) expected_outputs = {model.output_names[0]: model(inputs)} self._assert_all_close(expected_outputs, signature_outputs)
def test_trace_features_layer(self): columns = [feature_column_v2.numeric_column('x')] model = sequential.Sequential( [feature_column_v2.DenseFeatures(columns)]) model_input = {'x': constant_op.constant([[1.]])} model.predict(model_input, steps=1) fn = saving_utils.trace_model_call(model) self.assertAllClose({'output_1': [[1.]]}, fn({'x': [[1.]]})) columns = [feature_column_v2.numeric_column('x'), feature_column_v2.numeric_column('y')] model = sequential.Sequential( [feature_column_v2.DenseFeatures(columns)]) model_input = {'x': constant_op.constant([[1.]]), 'y': constant_op.constant([[2.]])} model.predict(model_input, steps=1) fn = saving_utils.trace_model_call(model) self.assertAllClose({'output_1': [[1., 2.]]}, fn({'x': [[1.]], 'y': [[2.]]}))
def from_keras_model(cls, model): """Creates a TFLiteConverter object from a Keras model. Args: model: tf.Keras.Model Returns: TFLiteConverter object. """ func = _saving_utils.trace_model_call(model) concrete_func = func.get_concrete_function() return cls([concrete_func])
def test_trace_model_outputs_after_fitting(self): input_dim = 5 if testing_utils.get_model_type() == 'functional' else None model = testing_utils.get_small_mlp(10, 3, input_dim) model.compile(optimizer='sgd', loss='mse') model.fit(x=np.random.random((8, 5)), y=np.random.random((8, 3)), epochs=2) inputs = array_ops.ones((8, 5)) fn = saving_utils.trace_model_call(model) signature_outputs = fn(inputs) expected_outputs = {model.output_names[0]: model(inputs)} self._assert_all_close(expected_outputs, signature_outputs)
def test_trace_multi_io_model_outputs(self): input_dim = 5 num_classes = 3 num_classes_b = 4 input_a = keras.layers.Input(shape=(input_dim,), name='input_a') input_b = keras.layers.Input(shape=(input_dim,), name='input_b') dense = keras.layers.Dense(num_classes, name='dense') dense2 = keras.layers.Dense(num_classes_b, name='dense2') dropout = keras.layers.Dropout(0.5, name='dropout') branch_a = [input_a, dense] branch_b = [input_b, dense, dense2, dropout] model = testing_utils.get_multi_io_model(branch_a, branch_b) input_a_np = np.random.random((10, input_dim)).astype(np.float32) input_b_np = np.random.random((10, input_dim)).astype(np.float32) if testing_utils.get_model_type() == 'subclass': with self.assertRaisesRegexp(ValueError, 'input shapes have not been set'): saving_utils.trace_model_call(model) model.compile(optimizer='sgd', loss='mse') model.fit(x=[np.random.random((8, input_dim)).astype(np.float32), np.random.random((8, input_dim)).astype(np.float32)], y=[np.random.random((8, num_classes)).astype(np.float32), np.random.random((8, num_classes_b)).astype(np.float32)], epochs=2) fn = saving_utils.trace_model_call(model) signature_outputs = fn([input_a_np, input_b_np]) outputs = model([input_a_np, input_b_np]) expected_outputs = {model.output_names[0]: outputs[0], model.output_names[1]: outputs[1]} self._assert_all_close(expected_outputs, signature_outputs)
def concrete_function_from_keras_model(model): input_signature = None if version.parse(tf.__version__) >= version.parse("2.1"): # If the model's call is not a `tf.function`, then we need to first get its # input signature from `model_input_signature` method. We can't directly # call `trace_model_call` because otherwise the batch dimension is set # to None. # Once we have better support for dynamic shapes, we can remove this. if not isinstance(model.call, def_function.Function): # Pass `keep_original_batch_size=True` will ensure that we get an input # signature including the batch dimension specified by the user. input_signature = saving_utils.model_input_signature( model, keep_original_batch_size=True ) func = saving_utils.trace_model_call(model, input_signature) return func.get_concrete_function()
def _from_tf_keras_model(keras_model): from tensorflow.python.keras.saving import saving_utils from tensorflow.python.framework.convert_to_constants import ( convert_variables_to_constants_v2, ) if not isinstance(keras_model, tf.keras.Model): keras_model = tf.keras.models.load_model(keras_model, None) tf.keras.backend.clear_session() tf.keras.backend.set_learning_phase(False) fn = saving_utils.trace_model_call(keras_model) cf = fn.get_concrete_function() try: frozen_fn = convert_variables_to_constants_v2(cf) return frozen_fn.graph.as_graph_def(add_shapes=True) except Exception: raise NotImplementedError("Unhandled tf.keras model format")
def from_keras(model_path, input_names, output_names): """Load keras model - experimental for now.""" from tensorflow.python import keras as _keras from tensorflow.python.eager import context from tensorflow.python.keras.saving import saving_utils as _saving_utils # Handles Keras when Eager mode is enabled. custom_objects = None with tf.device("/cpu:0"): if context.executing_eagerly(): _keras.backend.clear_session() _keras.backend.set_learning_phase(False) keras_model = _keras.models.load_model(model_path, custom_objects) function = _saving_utils.trace_model_call(keras_model) concrete_func = function.get_concrete_function() # allow to pass inputs and outputs from caller if we don't want all of them 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 ] frozen_graph = from_function(concrete_func, input_names, output_names) else: # Handles Keras when Eager mode is disabled. _keras.backend.clear_session() _keras.backend.set_learning_phase(False) keras_model = _keras.models.load_model(model_path, custom_objects) # allow to pass inputs and outputs from caller if we don't want all of them input_names = keras_model.inputs output_names = keras_model.outputs sess = _keras.backend.get_session() input_names = inputs_without_resource(sess, input_names) frozen_graph = freeze_session(sess, input_names=input_names, output_names=output_names) tf_reset_default_graph() with tf_session() as sess: frozen_graph = tf_optimize(input_names, output_names, frozen_graph) tf_reset_default_graph() return frozen_graph, input_names, output_names
def test_trace_model_outputs_after_fitting(self): input_dim = 5 if testing_utils.get_model_type( ) == 'functional' else None model = testing_utils.get_small_mlp(10, 3, input_dim) model.compile(optimizer='sgd', loss='mse', run_eagerly=testing_utils.should_run_eagerly()) model.fit(x=np.random.random((8, 5)), y=np.random.random((8, 3)), epochs=2) inputs = array_ops.ones((8, 5)) fn = saving_utils.trace_model_call(model) signature_outputs = fn(inputs) expected_outputs = {model.output_names[0]: model(inputs)} self._assert_all_close(expected_outputs, signature_outputs)
def extract_outputs_from_subclassing_model(model, output_dict, input_names, output_names): from tensorflow.python.keras.saving import saving_utils as _saving_utils from tensorflow.python.util import object_identity from ._graph_cvt import convert_variables_to_constants_v2 as _convert_to_constants function = _saving_utils.trace_model_call(model) concrete_func = function.get_concrete_function() output_names.extend([ts_.name for ts_ in concrete_func.outputs]) output_dict.update(build_layer_outputs(model, concrete_func.graph, concrete_func.outputs)) graph_def, converted_input_indices = _convert_to_constants( concrete_func, lower_control_flow=True) input_tensors = concrete_func.graph.internal_captures converted_inputs = object_identity.ObjectIdentitySet( [input_tensors[index] for index in converted_input_indices]) input_names.extend([ tensor.name for tensor in concrete_func.inputs if tensor not in converted_inputs]) with tf.Graph().as_default() as tf_graph: tf.import_graph_def(graph_def, name='') return tf_graph
def test_subclassed_model_with_input_signature(self): class Model(keras.Model): def __init__(self): super(Model, self).__init__() self.dense = keras.layers.Dense(3, name='dense') @def_function.function( input_signature=[[tensor_spec.TensorSpec([None, 5], dtypes.float32), tensor_spec.TensorSpec([None], dtypes.float32)]],) def call(self, inputs, *args): x, y = inputs return self.dense(x) + y model = Model() fn = saving_utils.trace_model_call(model) x = array_ops.ones((8, 5), dtype=dtypes.float32) y = array_ops.ones((3,), dtype=dtypes.float32) expected_outputs = {'output_1': model([x, y])} signature_outputs = fn([x, y]) self._assert_all_close(expected_outputs, signature_outputs)
def test_subclassed_model_with_input_signature(self): class Model(keras.Model): def __init__(self): super(Model, self).__init__() self.dense = keras.layers.Dense(3, name='dense') @def_function.function( input_signature=[[ tensor_spec.TensorSpec([None, 5], dtypes.float32), tensor_spec.TensorSpec([None], dtypes.float32) ]], ) def call(self, inputs, *args): x, y = inputs return self.dense(x) + y model = Model() fn = saving_utils.trace_model_call(model) x = array_ops.ones((8, 5), dtype=dtypes.float32) y = array_ops.ones((3, ), dtype=dtypes.float32) expected_outputs = {'output_1': model([x, y])} signature_outputs = fn([x, y]) self._assert_all_close(expected_outputs, signature_outputs)
def _graph_def_from_saved_model_or_keras_model(filename): """ Utility function that returns GraphDef object from the given SavedModel or HDF5 model. :param filename: TensorFlow SavedModel directory or Keras HDF5 model (.h5) file. :return: TensorFlow GraphDef object. """ try: import tensorflow as tf from tensorflow.python.keras.saving import saving_utils as _saving_utils from tensorflow.python.framework import convert_to_constants as _convert_to_constants model = tf.keras.models.load_model(filename) tf.keras.backend.set_learning_phase(False) func = _saving_utils.trace_model_call(model) concrete_func = func.get_concrete_function() # concrete_func = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY] frozen_func = _convert_to_constants.convert_variables_to_constants_v2(concrete_func) graph_def = frozen_func.graph.as_graph_def(add_shapes=True) except ImportError as e: raise ImportError('Failed to import TensorFlow utilities. {}.'.format(e)) except Exception as e: raise RuntimeError('Failed to load SavedModel or .h5 model. {}.'.format(e)) return graph_def
def _graph_def_from_saved_model_or_keras_model(filename): """ Utility function that returns GraphDef object from the given SavedModel or HDF5 model. :param filename: TensorFlow SavedModel directory or Keras HDF5 model (.h5) file. :return: TensorFlow GraphDef object. """ try: import tensorflow as tf from tensorflow.python.keras.saving import saving_utils as _saving_utils from tensorflow.python.framework import convert_to_constants as _convert_to_constants if filename.endswith('.h5'): model = tf.keras.models.load_model(filename) tf.keras.backend.set_learning_phase(False) func = _saving_utils.trace_model_call(model) concrete_func = func.get_concrete_function() else: model = tf.saved_model.load(filename) signatures = model.signatures if len(signatures) == 0: raise ValueError( 'Unable to load a model with no signatures provided.') if len(signatures) >= 2: raise ValueError( 'Unable to load a model with multiple signatures') concrete_func = list(signatures.values())[0] frozen_func = _convert_to_constants.convert_variables_to_constants_v2( concrete_func) graph_def = frozen_func.graph.as_graph_def(add_shapes=True) except ImportError as e: raise ImportError( 'Failed to import TensorFlow utilities. {}.'.format(e)) except ValueError as e: raise ValueError( 'Failed to load SavedModel or .h5 model. {}.'.format(e)) except Exception as e: raise RuntimeError( 'Failed to load SavedModel or .h5 model. {}.'.format(e)) return graph_def
def from_keras_model(cls, model): """Creates a TFLiteConverter object from a Keras model. Args: model: tf.Keras.Model Returns: TFLiteConverter object. """ input_signature = None # If the model's call is not a `tf.function`, then we need to first get its # input signature from `model_input_signature` method. We can't directly # call `trace_model_call` because otherwise the batch dimension is set # to None. # Once we have better support for dynamic shapes, we can remove this. if not isinstance(model.call, _def_function.Function): # Pass `keep_original_batch_size=True` will ensure that we get an input # signature including the batch dimension specified by the user. input_signature = _saving_utils.model_input_signature( model, keep_original_batch_size=True) func = _saving_utils.trace_model_call(model, input_signature) concrete_func = func.get_concrete_function() return cls([concrete_func])
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 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") if not input_signature: raise ValueError("from_keras requires input_signature") 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=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 _convert_tf_saved_model(output_dir, saved_model_dir=None, keras_model=None, signature_def='serving_default', saved_model_tags='serve', quantization_dtype_map=None, skip_op_check=False, strip_debug_ops=False, use_structured_outputs_names=False, weight_shard_size_bytes=1024 * 1024 * 4, control_flow_v2=False, experiments=False, metadata=None, frozen_graph_dir=None): """Take a SavedModel or KerasModel and convert to Tensorflow.js graph model. Args: output_dir: string The name of the output directory. The directory will consist of - a file named 'model.json' - possibly sharded binary weight files. saved_model_dir: string The saved model directory. : string The names of the output nodes, comma separated. keras_model: An in-memory Keras model object. signature_def: string Tagset of the SignatureDef to load. Defaults to 'serving_default'. saved_model_tags: tags of the GraphDef to load. Defaults to 'serve'. quantization_dtype_map: A mapping from dtype (`uint8`, `uint16`, `float16`) to weights names. The weight mapping supports wildcard substitution. skip_op_check: Bool whether to skip the op check. strip_debug_ops: Bool whether to strip debug ops. use_structured_outputs_names: Bool whether output of graph model will follow the structured_outputs format. weight_shard_size_bytes: Shard size (in bytes) of the weight files. The size of each weight file will be <= this value. control_flow_v2: Bool whether to enable control flow v2 ops. experiments: Bool enable experimental features. metadata: User defined metadata map. frozen_graph_dir: The directory to keep the intermediate frozen graph of model. """ if signature_def is None: signature_def = 'serving_default' if not tf.io.gfile.exists(output_dir): tf.io.gfile.makedirs(output_dir) output_graph = os.path.join(output_dir, common.ARTIFACT_MODEL_JSON_FILE_NAME) saved_model_tags_list = None if saved_model_tags: saved_model_tags_list = saved_model_tags.split(',') model = None concrete_func = None saved_model_sigature = None if saved_model_dir: saved_model_sigature = _find_signature(saved_model_dir, saved_model_tags, signature_def) model = _load_model(saved_model_dir, saved_model_tags_list) _check_signature_in_model(model, signature_def) concrete_func = model.signatures[signature_def] elif keras_model: model = keras_model input_signature = None # If the model's call is not a `tf.function`, then we need to first get its # input signature from `model_input_signature` method. We can't directly # call `trace_model_call` because otherwise the batch dimension is set # to None. if not isinstance(model.call, def_function.Function): # Pass `keep_original_batch_size=True` will ensure that we get an input # signature including the batch dimension specified by the user. input_signature = model_input_signature( model, keep_original_batch_size=True) func = trace_model_call(model, input_signature) concrete_func = func.get_concrete_function() else: raise Exception( 'Provide either a saved model or keras model to convert.') output_node_names = [] for output_tensor in concrete_func.outputs: output_node_names.append(output_tensor.name.split(':')[0]) num_outputs = len(output_node_names) structured_outputs = concrete_func.structured_outputs if use_structured_outputs_names and structured_outputs is not None: if not isinstance(structured_outputs, dict): raise Exception('Converter only supports dict structured_outputs.') # As per tensorflow/python/util/nest.py: "If `structure` is or contains a # dict instance, the keys will be sorted to pack the flat sequence # in deterministic order." sorted_keys = sorted(structured_outputs.keys()) # Check if structure is a simple dictionary. # We don't support anything more complex due to the GraphModel.predict # function return type in typescript. test_sequence = list(range(num_outputs)) actual_structure = tf.nest.pack_sequence_as(structured_outputs, test_sequence, True) expected_structure = dict(zip(sorted_keys, test_sequence)) if actual_structure != expected_structure: raise Exception( 'Converter only supports structured_outputs of form ' '{"key1": value1, "key2":value2 ... })') metadata = metadata or {} metadata[common.STRUCTURED_OUTPUTS_KEYS_KEY] = sorted_keys # TensorFlow doesn't encode the saved model version in the graph in a # reliable way. Try to freeze the graph using V2 utils. If that fails, freeze # the graph using V1 utils. frozen_initializer_graph = None try: frozen_graph = _freeze_saved_model_v2(concrete_func, control_flow_v2) except BaseException: if saved_model_dir: (frozen_graph, frozen_initializer_graph) = _freeze_saved_model_v1( saved_model_dir, saved_model_tags_list, output_node_names) else: print('Can not freeze saved model v1.') return if frozen_graph_dir: output_graph = os.path.join(frozen_graph_dir, common.ARTIFACT_MODEL_JSON_FILE_NAME) frozen_file = output_graph + '.frozen' with tf.compat.v1.gfile.GFile(frozen_file, 'wb') as f: f.write(frozen_graph.as_graph_def().SerializeToString()) inputs = [x for x in concrete_func.inputs if not x.dtype == 'resource'] signature = _build_signature_def(frozen_graph, inputs, concrete_func.outputs, saved_model_sigature) define_transform_graph_func() tf_version = None try: tf_version = model.tensorflow_version except: # pylint: disable=W0702 # keras model does not have tensorflow_version, hard code to the latest # tensorflow version. tf_version = tf.__version__ optimize_graph(frozen_graph, signature, output_graph, tf_version, quantization_dtype_map=quantization_dtype_map, skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops, weight_shard_size_bytes=weight_shard_size_bytes, experiments=experiments, initializer_graph=frozen_initializer_graph, metadata=metadata)
def export_saved_model(model, saved_model_path, custom_objects=None, as_text=False, input_signature=None, serving_only=False): """Exports a `tf.keras.Model` as a Tensorflow SavedModel. Note that at this time, subclassed models can only be saved using `serving_only=True`. The exported `SavedModel` is a standalone serialization of Tensorflow objects, and is supported by TF language APIs and the Tensorflow Serving system. To load the model, use the function `tf.keras.experimental.load_from_saved_model`. The `SavedModel` contains: 1. a checkpoint containing the model weights. 2. a `SavedModel` proto containing the Tensorflow backend graph. Separate graphs are saved for prediction (serving), train, and evaluation. If the model has not been compiled, then only the graph computing predictions will be exported. 3. the model's json config. If the model is subclassed, this will only be included if the model's `get_config()` method is overwritten. Example: ```python import tensorflow as tf # Create a tf.keras model. model = tf.keras.Sequential() model.add(tf.keras.layers.Dense(1, input_shape=[10])) model.summary() # Save the tf.keras model in the SavedModel format. path = '/tmp/simple_keras_model' tf.keras.experimental.export_saved_model(model, path) # Load the saved keras model back. new_model = tf.keras.experimental.load_from_saved_model(path) new_model.summary() ``` Args: model: A `tf.keras.Model` to be saved. If the model is subclassed, the flag `serving_only` must be set to True. saved_model_path: a string specifying the path to the SavedModel directory. custom_objects: Optional dictionary mapping string names to custom classes or functions (e.g. custom loss functions). as_text: bool, `False` by default. Whether to write the `SavedModel` proto in text format. Currently unavailable in serving-only mode. input_signature: A possibly nested sequence of `tf.TensorSpec` objects, used to specify the expected model inputs. See `tf.function` for more details. serving_only: bool, `False` by default. When this is true, only the prediction graph is saved. Raises: NotImplementedError: If the model is a subclassed model, and serving_only is False. ValueError: If the input signature cannot be inferred from the model. AssertionError: If the SavedModel directory already exists and isn't empty. """ warnings.warn('`tf.keras.experimental.export_saved_model` is deprecated' 'and will be removed in a future version. ' 'Please use `model.save(..., save_format="tf")` or ' '`tf.keras.models.save_model(..., save_format="tf")`.') if serving_only: save_lib.save( model, saved_model_path, signatures=saving_utils.trace_model_call(model, input_signature)) else: _save_v1_format(model, saved_model_path, custom_objects, as_text, input_signature) try: _export_model_json(model, saved_model_path) except NotImplementedError: logging.warning('Skipped saving model JSON, subclassed model does not have ' 'get_config() defined.')
def export(model, saved_model_path, custom_objects=None, as_text=None, input_signature=None, serving_only=False): """Saves a `tf.keras.Model` into Tensorflow SavedModel format. `save_model` generates new files/folders under the `saved_model_path` folder: 1) a checkpoint containing the model weights. 2) a saved_model.pb file containing the model's MetaGraphs. The prediction graph is always exported. The evaluaton and training graphs are exported if the following conditions are met: - Evaluation: model loss is defined. - Training: model is compiled with an optimizer defined under `tf.train`. This is because `tf.keras.optimizers.Optimizer` instances cannot be saved to checkpoints. 3) Model's json configuration, if model.get_config() has been implemented. This file can be used to reload the model using tf.keras.models.model_from_json(). Note that if any custom objects were used, they should be passed to the `custom_object` argument when loading the model. Model limitations: - Sequential and functional models can always be saved. - Subclassed models can only be saved when `serving_only=True`. This is due to the current implementation copying the model in order to export the training and evaluation graphs. Because the topology of subclassed models cannot be determined, the subclassed models cannot be cloned. Subclassed models will be entirely exportable in the future. Note that each mode is exported in separate graphs, so different modes do not share variables. To use the train graph with evaluation or prediction graphs, create a new checkpoint if variable values have been updated. Example: ```python import tensorflow as tf # Create a tf.keras model. model = tf.keras.Sequential() model.add(tf.keras.layers.Dense(1, input_shape=[10])) model.summary() # Save the tf.keras model in the SavedModel format. saved_to_path = tf.keras.experimental.export( model, '/tmp/my_simple_tf_keras_saved_model') # Load the saved keras model back. model_prime = tf.keras.experimental.load_from_saved_model(saved_to_path) model_prime.summary() ``` Args: model: A `tf.keras.Model` to be saved. If the model is subclassed, the flag `serving_only` must be set to True. saved_model_path: a string specifying the path to the SavedModel directory. The SavedModel will be saved to a timestamped folder created within this directory. custom_objects: Optional dictionary mapping string names to custom classes or functions (e.g. custom loss functions). as_text: whether to write the `SavedModel` proto in text format. Currently unavailable in serving-only mode. input_signature: A possibly nested sequence of `tf.TensorSpec` objects, used to specify the expected model inputs. `input_signature`'s nested structure should match the expected nested structure of the inputs to the model. If this is not set, this function will attempt to infer the input shapes and dtypes from the model. Note that if the model is subclassed, the tensor inputs to the call function should be nested in the first argument (this is a general requirement for using subclassed models with Keras functions .fit(), .predict(), etc.). serving_only: Export only the outputs produced from calling the model in predict mode. The losses, optimizer, and other training configurations are not saved. If the SavedModel will only be used for serving (rather than retraining), or if the model is subclassed, this can be set to True. Returns: String path to the SavedModel folder, a subdirectory of `saved_model_path`. Raises: NotImplementedError: If the model is a subclassed model, and serving_only is False. ValueError: If the input signature cannot be inferred from the model. """ export_dir = model_utils.get_timestamped_export_dir(saved_model_path) if serving_only: save_lib.save(model, export_dir, signatures=saving_utils.trace_model_call( model, input_signature)) else: _save_v1_format(model, export_dir, custom_objects, as_text, input_signature) try: _export_model_json(model, export_dir) except NotImplementedError: logging.warning( 'Skipped saving model JSON, subclassed model does not have ' 'get_config() defined.') return export_dir
def dump(cls, path, obj, variables_to_constants=False, output_names=None, *args, **kwargs): """ Extracts a TensorFlow graph from an object *obj* and saves it at *path*. The graph is optionally transformed into a simpler representation with all its variables converted to constants when *variables_to_constants* is *True*. The saved file contains the graph as a protobuf. The accepted types of *obj* greatly depend on the available API versions. When the v1 API is found (which is also the case when ``tf.compat.v1`` is available in v2), ``Graph``, ``GraphDef`` and ``Session`` objects are accepted. However, when *variables_to_constants* is *True*, *obj* must be a session and *output_names* should refer to names of operations whose subgraphs are extracted (usually just one). For TensorFlow v2, *obj* can also be a compiled keras model, or either a polymorphic or concrete function as returned by ``tf.function``. Polymorphic functions either must have a defined input signature (``tf.function(input_signature=(...,))``) or they must accept no arguments in the first place. See the TensorFlow documentation on `concrete functions <https://www.tensorflow.org/guide/concrete_function>`__ for more info. *args* and *kwargs* are forwarded to ``tf.train.write_graph`` (v1) or ``tf.io.write_graph`` (v2). """ tf, tf1, tf_version = cls.import_tf() path = get_path(path) graph_dir, graph_name = os.path.split(path) # default as_text value kwargs.setdefault("as_text", path.endswith((".pbtxt", ".pb.txt"))) # convert keras models and polymorphic functions to concrete functions, v2 only if tf_version[0] != "1": from tensorflow.python.keras.saving import saving_utils from tensorflow.python.eager.def_function import Function from tensorflow.python.eager.function import ConcreteFunction if isinstance(obj, tf.keras.Model): learning_phase_orig = tf.keras.backend.learning_phase() tf.keras.backend.set_learning_phase(False) model_func = saving_utils.trace_model_call(obj) if model_func.function_spec.arg_names and not model_func.input_signature: raise ValueError("when obj is a keras model callable accepting arguments, its " "input signature must be frozen by building the model") obj = model_func.get_concrete_function() tf.keras.backend.set_learning_phase(learning_phase_orig) elif isinstance(obj, Function): if obj.function_spec.arg_names and not obj.input_signature: raise ValueError("when obj is a polymorphic function accepting arguments, its " "input signature must be frozen") obj = obj.get_concrete_function() # convert variables to constants if variables_to_constants: if tf1 and isinstance(obj, tf1.Session): if not output_names: raise ValueError("when variables_to_constants is true, output_names must " "contain operations to export, got '{}' instead".format(output_names)) obj = tf1.graph_util.convert_variables_to_constants(obj, obj.graph.as_graph_def(), output_names) elif tf_version[0] != "1": from tensorflow.python.framework import convert_to_constants if not isinstance(obj, ConcreteFunction): raise TypeError("when variables_to_constants is true, obj must be a concrete " "or polymorphic function, got '{}' instead".format(obj)) obj = convert_to_constants.convert_variables_to_constants_v2(obj) else: raise TypeError("cannot convert variables to constants for object '{}', type not " "understood for TensorFlow version {}".format(obj, tf.__version__)) # extract the graph if tf1 and isinstance(obj, tf1.Session): graph = obj.graph elif tf_version[0] != "1" and isinstance(obj, ConcreteFunction): graph = obj.graph else: graph = obj # write it if tf_version[0] == "1": return tf1.train.write_graph(graph, graph_dir, graph_name, *args, **kwargs) else: return tf.io.write_graph(graph, graph_dir, graph_name, *args, **kwargs)
def from_keras_model_file(cls, model_file, input_arrays=None, input_shapes=None, output_arrays=None, custom_objects=None): """Creates a TFLiteConverter class from a tf.keras model file. Args: model_file: Full filepath of HDF5 file containing the tf.keras model. input_arrays: List of input tensors to freeze graph with. Uses input arrays from SignatureDef when none are provided. (default None) input_shapes: Dict of strings representing input tensor names to list of integers representing input shapes (e.g., {"foo" : [1, 16, 16, 3]}). Automatically determined when input shapes is None (e.g., {"foo" : None}). (default None) output_arrays: List of output tensors to freeze graph with. Uses output arrays from SignatureDef when none are provided. (default None) custom_objects: Dict mapping names (strings) to custom classes or functions to be considered during model deserialization. (default None) Returns: TFLiteConverter class. """ # Handles Keras when Eager mode is enabled. if context.executing_eagerly(): if input_arrays or output_arrays: raise ValueError("`input_arrays` and `output_arrays` are unsupported " "with Eager mode. If your model requires any of these " "parameters, please use disable_eager_execution().") _keras.backend.set_learning_phase(False) keras_model = _keras.models.load_model(model_file, custom_objects) function = _saving_utils.trace_model_call(keras_model) concrete_func = function.get_concrete_function() frozen_func = _convert_to_constants.convert_variables_to_constants_v2( concrete_func) _set_tensor_shapes(frozen_func.inputs, input_shapes) return cls(frozen_func.graph.as_graph_def(), frozen_func.inputs, frozen_func.outputs) # Handles Keras when Eager mode is disabled. _keras.backend.clear_session() _keras.backend.set_learning_phase(False) keras_model = _keras.models.load_model(model_file, custom_objects) sess = _keras.backend.get_session() # Get input and output tensors. if input_arrays: input_tensors = _get_tensors_from_tensor_names(sess.graph, input_arrays) else: input_tensors = keras_model.inputs if output_arrays: output_tensors = _get_tensors_from_tensor_names(sess.graph, output_arrays) else: output_tensors = keras_model.outputs _set_tensor_shapes(input_tensors, input_shapes) graph_def = _freeze_graph(sess, input_tensors, output_tensors) return cls(graph_def, input_tensors, output_tensors)
def export( model, saved_model_path, custom_objects=None, as_text=None, input_signature=None, serving_only=False): """Saves a `tf.keras.Model` into Tensorflow SavedModel format. `save_model` generates new files/folders under the `saved_model_path` folder: 1) a checkpoint containing the model weights. 2) a saved_model.pb file containing the model's MetaGraphs. The prediction graph is always exported. The evaluaton and training graphs are exported if the following conditions are met: - Evaluation: model loss is defined. - Training: model is compiled with an optimizer defined under `tf.train`. This is because `tf.keras.optimizers.Optimizer` instances cannot be saved to checkpoints. 3) Model's json configuration, if model.get_config() has been implemented. This file can be used to reload the model using tf.keras.models.model_from_json(). Note that if any custom objects were used, they should be passed to the `custom_object` argument when loading the model. Model limitations: - Sequential and functional models can always be saved. - Subclassed models can only be saved when `serving_only=True`. This is due to the current implementation copying the model in order to export the training and evaluation graphs. Because the topology of subclassed models cannot be determined, the subclassed models cannot be cloned. Subclassed models will be entirely exportable in the future. Note that each mode is exported in separate graphs, so different modes do not share variables. To use the train graph with evaluation or prediction graphs, create a new checkpoint if variable values have been updated. Example: ```python import tensorflow as tf # Create a tf.keras model. model = tf.keras.Sequential() model.add(tf.keras.layers.Dense(1, input_shape=[10])) model.summary() # Save the tf.keras model in the SavedModel format. saved_to_path = tf.keras.experimental.export( model, '/tmp/my_simple_tf_keras_saved_model') # Load the saved keras model back. model_prime = tf.keras.experimental.load_from_saved_model(saved_to_path) model_prime.summary() ``` Args: model: A `tf.keras.Model` to be saved. If the model is subclassed, the flag `serving_only` must be set to True. saved_model_path: a string specifying the path to the SavedModel directory. The SavedModel will be saved to a timestamped folder created within this directory. custom_objects: Optional dictionary mapping string names to custom classes or functions (e.g. custom loss functions). as_text: whether to write the `SavedModel` proto in text format. Currently unavailable in serving-only mode. input_signature: A possibly nested sequence of `tf.TensorSpec` objects, used to specify the expected model inputs. `input_signature`'s nested structure should match the expected nested structure of the inputs to the model. If this is not set, this function will attempt to infer the input shapes and dtypes from the model. Note that if the model is subclassed, the tensor inputs to the call function should be nested in the first argument (this is a general requirement for using subclassed models with Keras functions .fit(), .predict(), etc.). serving_only: Export only the outputs produced from calling the model in predict mode. The losses, optimizer, and other training configurations are not saved. If the SavedModel will only be used for serving (rather than retraining), or if the model is subclassed, this can be set to True. Returns: String path to the SavedModel folder, a subdirectory of `saved_model_path`. Raises: NotImplementedError: If the model is a subclassed model, and serving_only is False. ValueError: If the input signature cannot be inferred from the model. """ export_dir = model_utils.get_timestamped_export_dir(saved_model_path) if serving_only: save_lib.save( model, export_dir, signatures=saving_utils.trace_model_call(model, input_signature)) else: _save_v1_format(model, export_dir, custom_objects, as_text, input_signature) try: _export_model_json(model, export_dir) except NotImplementedError: logging.warning('Skipped saving model JSON, subclassed model does not have ' 'get_config() defined.') return export_dir
def from_keras_model_file(cls, model_file, input_arrays=None, input_shapes=None, output_arrays=None, custom_objects=None): """Creates a TFLiteConverter class from a tf.keras model file. Args: model_file: Full filepath of HDF5 file containing the tf.keras model. input_arrays: List of input tensors to freeze graph with. Uses input arrays from SignatureDef when none are provided. (default None) input_shapes: Dict of strings representing input tensor names to list of integers representing input shapes (e.g., {"foo" : [1, 16, 16, 3]}). Automatically determined when input shapes is None (e.g., {"foo" : None}). (default None) output_arrays: List of output tensors to freeze graph with. Uses output arrays from SignatureDef when none are provided. (default None) custom_objects: Dict mapping names (strings) to custom classes or functions to be considered during model deserialization. (default None) Returns: TFLiteConverter class. """ # Handles Keras when Eager mode is enabled. if context.executing_eagerly(): if input_arrays or output_arrays: raise ValueError( "`input_arrays` and `output_arrays` are unsupported " "with Eager mode. If your model requires any of these " "parameters, please use disable_eager_execution().") _keras.backend.set_learning_phase(False) keras_model = _keras.models.load_model(model_file, custom_objects) function = _saving_utils.trace_model_call(keras_model) concrete_func = function.get_concrete_function() frozen_func = _convert_to_constants.convert_variables_to_constants_v2( concrete_func, lower_control_flow=False) _set_tensor_shapes(frozen_func.inputs, input_shapes) return cls(frozen_func.graph.as_graph_def(), frozen_func.inputs, frozen_func.outputs, experimental_debug_info_func=_build_debug_info_func( frozen_func.graph)) # Handles Keras when Eager mode is disabled. _keras.backend.clear_session() _keras.backend.set_learning_phase(False) keras_model = _keras.models.load_model(model_file, custom_objects) sess = _keras.backend.get_session() # Get input and output tensors. if input_arrays: input_tensors = _get_tensors_from_tensor_names( sess.graph, input_arrays) else: input_tensors = keras_model.inputs if output_arrays: output_tensors = _get_tensors_from_tensor_names( sess.graph, output_arrays) else: output_tensors = keras_model.outputs _set_tensor_shapes(input_tensors, input_shapes) graph_def = _freeze_graph(sess, input_tensors, output_tensors) return cls(graph_def, input_tensors, output_tensors, experimental_debug_info_func=_build_debug_info_func( sess.graph))
def _convert_tf_saved_model(output_dir, saved_model_dir=None, keras_model=None, signature_def='serving_default', saved_model_tags='serve', quantization_dtype_map=None, skip_op_check=False, strip_debug_ops=False, weight_shard_size_bytes=1024 * 1024 * 4, control_flow_v2=False, experiments=False, metadata=None): """Take a SavedModel or KerasModel and convert to Tensorflow.js graph model. Args: output_dir: string The name of the output directory. The directory will consist of - a file named 'model.json' - possibly sharded binary weight files. saved_model_dir: string The saved model directory. : string The names of the output nodes, comma separated. keras_model: An in-memory Keras model object. signature_def: string Tagset of the SignatureDef to load. Defaults to 'serving_default'. saved_model_tags: tags of the GraphDef to load. Defaults to 'serve'. quantization_dtype_map: A mapping from dtype (`uint8`, `uint16`, `float16`) to weights names. The weight mapping supports wildcard substitution. skip_op_check: Bool whether to skip the op check. strip_debug_ops: Bool whether to strip debug ops. weight_shard_size_bytes: Shard size (in bytes) of the weight files. The size of each weight file will be <= this value. control_flow_v2: Bool whether to enable control flow v2 ops. experiments: Bool enable experimental features. metadata: User defined metadata map. """ if signature_def is None: signature_def = 'serving_default' if not tf.io.gfile.exists(output_dir): tf.io.gfile.makedirs(output_dir) output_graph = os.path.join(output_dir, common.ARTIFACT_MODEL_JSON_FILE_NAME) saved_model_tags_list = None if saved_model_tags: saved_model_tags_list = saved_model_tags.split(',') model = None concrete_func = None saved_model_sigature = None if saved_model_dir: saved_model_sigature = _find_signature(saved_model_dir, saved_model_tags, signature_def) model = _load_model(saved_model_dir, saved_model_tags_list) _check_signature_in_model(model, signature_def) concrete_func = model.signatures[signature_def] elif keras_model: model = keras_model input_signature = None # If the model's call is not a `tf.function`, then we need to first get its # input signature from `model_input_signature` method. We can't directly # call `trace_model_call` because otherwise the batch dimension is set # to None. if not isinstance(model.call, def_function.Function): # Pass `keep_original_batch_size=True` will ensure that we get an input # signature including the batch dimension specified by the user. input_signature = model_input_signature( model, keep_original_batch_size=True) func = trace_model_call(model, input_signature) concrete_func = func.get_concrete_function() else: raise Exception( 'Provide either a saved model or keras model to convert.') output_node_names = [] for output_tensor in concrete_func.outputs: output_node_names.append(output_tensor.name.split(':')[0]) # TensorFlow doesn't encode the saved model version in the graph in a # reliable way. Try to freeze the graph using V2 utils. If that fails, freeze # the graph using V1 utils. frozen_initializer_graph = None try: frozen_graph = _freeze_saved_model_v2(concrete_func, control_flow_v2) except BaseException: if saved_model_dir: (frozen_graph, frozen_initializer_graph) = _freeze_saved_model_v1( saved_model_dir, saved_model_tags_list, output_node_names) else: print('Can not freeze saved model v1.') return inputs = [x for x in concrete_func.inputs if not x.dtype == 'resource'] signature = _build_signature_def(frozen_graph, inputs, concrete_func.outputs, saved_model_sigature) # Check if the TransformGraph is available to be imported, this package is # available in g3 but not in oss version of TensorFlow. transform_graph_available = True try: from tensorflow.tools.graph_transforms import TransformGraph # pylint: disable=C0415 except: # pylint: disable=W0702 transform_graph_available = False # Define the strip graph functions when TransformGraph is available, this will # strip the unused nodes from the graph. if transform_graph_available: def _strip_unused_nodes(frozen_graph, concrete_func, output_node_names): # Find the names of the input nodes needed to extract the minimal # inference graph. This is particularly useful for cases when the concrete # function contains nodes that do not contribute the inference computation # defined by the input/output pair. This would also eliminate op # unsupported error caused by nodes outside of the minial infrerence # graph. input_node_names = [] input_tensors = {} for input_tensor in concrete_func.inputs: if input_tensor.dtype != 'resource': op_name = input_tensor.name.split(':')[0] # The graph freezing may turn the original inputs into constants, or # remove them from the graph, so we need to ignore those. try: op = frozen_graph.get_operation_by_name(op_name) if op.type != 'Const': input_node_names.append(op_name) input_tensors[op_name] = input_tensor except KeyError: # The original input was removed when the graph was frozen. continue graph_transformations = ['strip_unused_nodes'] stripped_graph_def = TransformGraph(frozen_graph.as_graph_def(), input_node_names, output_node_names, graph_transformations) # The transform graph library cannot support input nodes that has dynamic # shape, this code will update the dtype and shape based on the # input tensor manually. for node in stripped_graph_def.node: if node.name in input_tensors: if node.attr['shape'] and node.attr['shape'].shape: node.attr['shape'].shape.CopyFrom( input_tensors[node.name].shape.as_proto()) if node.attr['dtype'] and node.attr['dtype'].type: node.attr['dtype'].type = input_tensors[ node.name].dtype.as_datatype_enum with tf.Graph().as_default() as stripped_graph: tf.import_graph_def(stripped_graph_def, name='') return stripped_graph frozen_graph = _strip_unused_nodes(frozen_graph, concrete_func, output_node_names) version = None try: version = model.tensorflow_version except: # pylint: disable=W0702 # keras model does not have tensorflow_version, hard code to the latest # tensorflow version. version = tf.__version__ optimize_graph(frozen_graph, signature, output_graph, version, quantization_dtype_map=quantization_dtype_map, skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops, weight_shard_size_bytes=weight_shard_size_bytes, experiments=experiments, initializer_graph=frozen_initializer_graph, metadata=metadata)
def export_saved_model(model, saved_model_path, custom_objects=None, as_text=False, input_signature=None, serving_only=False): """Exports a `tf.keras.Model` as a Tensorflow SavedModel. Note that at this time, subclassed models can only be saved using `serving_only=True`. The exported `SavedModel` is a standalone serialization of Tensorflow objects, and is supported by TF language APIs and the Tensorflow Serving system. To load the model, use the function `tf.keras.experimental.load_from_saved_model`. The `SavedModel` contains: 1. a checkpoint containing the model weights. 2. a `SavedModel` proto containing the Tensorflow backend graph. Separate graphs are saved for prediction (serving), train, and evaluation. If the model has not been compiled, then only the graph computing predictions will be exported. 3. the model's json config. If the model is subclassed, this will only be included if the model's `get_config()` method is overwritten. Example: ```python import tensorflow as tf # Create a tf.keras model. model = tf.keras.Sequential() model.add(tf.keras.layers.Dense(1, input_shape=[10])) model.summary() # Save the tf.keras model in the SavedModel format. path = '/tmp/simple_keras_model' tf.keras.experimental.export_saved_model(model, path) # Load the saved keras model back. new_model = tf.keras.experimental.load_from_saved_model(path) new_model.summary() ``` Args: model: A `tf.keras.Model` to be saved. If the model is subclassed, the flag `serving_only` must be set to True. saved_model_path: a string specifying the path to the SavedModel directory. custom_objects: Optional dictionary mapping string names to custom classes or functions (e.g. custom loss functions). as_text: bool, `False` by default. Whether to write the `SavedModel` proto in text format. Currently unavailable in serving-only mode. input_signature: A possibly nested sequence of `tf.TensorSpec` objects, used to specify the expected model inputs. See `tf.function` for more details. serving_only: bool, `False` by default. When this is true, only the prediction graph is saved. Raises: NotImplementedError: If the model is a subclassed model, and serving_only is False. ValueError: If the input signature cannot be inferred from the model. AssertionError: If the SavedModel directory already exists and isn't empty. """ if serving_only: save_lib.save( model, saved_model_path, signatures=saving_utils.trace_model_call(model, input_signature)) else: _save_v1_format(model, saved_model_path, custom_objects, as_text, input_signature) try: _export_model_json(model, saved_model_path) except NotImplementedError: logging.warning('Skipped saving model JSON, subclassed model does not have ' 'get_config() defined.')
def default_save_signature(layer): original_losses = _reset_layer_losses(layer) fn = saving_utils.trace_model_call(layer) fn.get_concrete_function() _restore_layer_losses(original_losses) return fn
def _convert_tf_saved_model(output_dir, saved_model_dir=None, keras_model=None, signature_def='serving_default', saved_model_tags='serve', quantization_dtype_map=None, skip_op_check=False, strip_debug_ops=False, weight_shard_size_bytes=1024 * 1024 * 4, control_flow_v2=False, experiments=False, metadata=None): """Take a SavedModel or KerasModel and convert to Tensorflow.js graph model. Args: output_dir: string The name of the output directory. The directory will consist of - a file named 'model.json' - possibly sharded binary weight files. saved_model_dir: string The saved model directory. : string The names of the output nodes, comma separated. keras_model: An in-memory Keras model object. signature_def: string Tagset of the SignatureDef to load. Defaults to 'serving_default'. saved_model_tags: tags of the GraphDef to load. Defaults to 'serve'. quantization_dtype_map: A mapping from dtype (`uint8`, `uint16`, `float16`) to weights names. The weight mapping supports wildcard substitution. skip_op_check: Bool whether to skip the op check. strip_debug_ops: Bool whether to strip debug ops. weight_shard_size_bytes: Shard size (in bytes) of the weight files. The size of each weight file will be <= this value. control_flow_v2: Bool whether to enable control flow v2 ops. experiments: Bool enable experimental features. metadata: User defined metadata map. """ if signature_def is None: signature_def = 'serving_default' if not tf.io.gfile.exists(output_dir): tf.io.gfile.makedirs(output_dir) output_graph = os.path.join( output_dir, common.ARTIFACT_MODEL_JSON_FILE_NAME) saved_model_tags_list = None if saved_model_tags: saved_model_tags_list = saved_model_tags.split(',') model = None concrete_func = None saved_model_sigature = None if saved_model_dir: saved_model_sigature = _find_signature(saved_model_dir, saved_model_tags, signature_def) model = _load_model(saved_model_dir, saved_model_tags_list) _check_signature_in_model(model, signature_def) concrete_func = model.signatures[signature_def] elif keras_model: model = keras_model input_signature = None # If the model's call is not a `tf.function`, then we need to first get its # input signature from `model_input_signature` method. We can't directly # call `trace_model_call` because otherwise the batch dimension is set # to None. if not isinstance(model.call, def_function.Function): # Pass `keep_original_batch_size=True` will ensure that we get an input # signature including the batch dimension specified by the user. input_signature = model_input_signature( model, keep_original_batch_size=True) func = trace_model_call(model, input_signature) concrete_func = func.get_concrete_function() else: raise Exception('Provide either a saved model or keras model to convert.') output_node_names = [] for output_tensor in concrete_func.outputs: output_node_names.append(output_tensor.name.split(':')[0]) # TensorFlow doesn't encode the saved model version in the graph in a # reliable way. Try to freeze the graph using V2 utils. If that fails, freeze # the graph using V1 utils. frozen_initializer_graph = None try: frozen_graph = _freeze_saved_model_v2(concrete_func, control_flow_v2) except BaseException: if saved_model_dir: (frozen_graph, frozen_initializer_graph) = _freeze_saved_model_v1(saved_model_dir, saved_model_tags_list, output_node_names) else: print('Can not freeze saved model v1.') return inputs = [x for x in concrete_func.inputs if not x.dtype == 'resource'] signature = _build_signature_def( frozen_graph, inputs, concrete_func.outputs, saved_model_sigature) define_transform_graph_func() version = None try: version = model.tensorflow_version except: # pylint: disable=W0702 # keras model does not have tensorflow_version, hard code to the latest # tensorflow version. version = tf.__version__ optimize_graph(frozen_graph, signature, output_graph, version, quantization_dtype_map=quantization_dtype_map, skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops, weight_shard_size_bytes=weight_shard_size_bytes, experiments=experiments, initializer_graph=frozen_initializer_graph, metadata=metadata)