def _save_v1_format(model, path, custom_objects, as_text, input_signature): """Exports model to v1 SavedModel format.""" from tensorflow.python.keras.engine import sequential # pylint: disable=g-import-not-at-top if not model._is_graph_network: if isinstance(model, sequential.Sequential): # If input shape is not directly set in the model, the exported model # will infer the expected shapes of the input from the model. if not model.built and input_signature is None: raise ValueError( 'Sequential model\'s input shape is unknown. Please build the ' 'model, or use the input_signature argument to specify the ' 'model inputs.') else: raise NotImplementedError( 'Subclassed models can only be exported for serving. Please set ' 'argument serving_only=True.') builder = saved_model_builder._SavedModelBuilder(path) # Manually save variables to export them in an object-based checkpoint. This # skips the `builder.add_meta_graph_and_variables()` step, which saves a # named-based checkpoint. # TODO(b/113134168): Add fn to Builder to save with object-based saver. # TODO(b/113178242): This should only export the model json structure. Only # one save is needed once the weights can be copied from the model to clone. checkpoint_path = _export_model_variables(model, path) # Export each mode. Use ModeKeys enums defined for `Estimator` to ensure that # Keras models and `Estimator`s are exported with the same format. # Every time a mode is exported, the code checks to see if new variables have # been created (e.g. optimizer slot variables). If that is the case, the # checkpoint is re-saved to include the new variables. export_args = { 'builder': builder, 'model': model, 'custom_objects': custom_objects, 'checkpoint_path': checkpoint_path, 'input_signature': input_signature } has_saved_vars = False if model.optimizer: # TODO(kathywu): Verify this works with v2 optimizer. if isinstance(model.optimizer, optimizers.TFOptimizer): _export_mode(mode_keys.ModeKeys.TRAIN, has_saved_vars, **export_args) has_saved_vars = True _export_mode(mode_keys.ModeKeys.TEST, has_saved_vars, **export_args) else: logging.warning( 'Model was compiled with an optimizer, but the optimizer is not from ' '`tf.train` (e.g. `tf.train.AdagradOptimizer`). Only the serving ' 'graph was exported. The train and evaluate graphs were not added to ' 'the SavedModel.') _export_mode(mode_keys.ModeKeys.PREDICT, has_saved_vars, **export_args) builder.save(as_text)
def save_keras_model(model, saved_model_path, custom_objects=None, as_text=None): """Save a `tf.keras.Model` into Tensorflow SavedModel format. `save_model` generates new files/folders under the `saved_model_path` folder: 1) an asset folder containing the json string of the model's configuration (topology). 2) a checkpoint containing the model weights. 3) 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. Model Requirements: - Model must be a sequential model or functional model. Subclassed models can not be saved via this function, unless you provide an implementation for get_config() and from_config(). - All variables must be saveable by the model. In general, this condition is met through the use of layers defined in the keras library. However, there is currently a bug with variables created in Lambda layer functions not being saved correctly (see https://github.com/keras-team/keras/issues/9740). 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.contrib.saved_model.save_keras_model( model, '/tmp/my_simple_tf_keras_saved_model') # Load the saved keras model back. model_prime = tf.contrib.saved_model.load_keras_model(saved_to_path) model_prime.summary() ``` Args: model: A `tf.keras.Model` to be saved. 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. Returns: String path to the SavedModel folder, a subdirectory of `saved_model_path`. Raises: NotImplementedError: If the model is a subclassed model. ValueError: If a Sequential model does not have input shapes defined by the user, and is not built. """ if not model._is_graph_network: if isinstance(model, sequential.Sequential): # If input shape is not directly set in the model, the exported model # will assume that the inputs have the same shape as the shape the model # was built model with. if not model.built: raise ValueError( 'Sequential model must be built before it can be exported.' ) else: raise NotImplementedError( 'Exporting subclassed models is not yet supported.') export_dir = export_helpers.get_timestamped_export_dir(saved_model_path) temp_export_dir = export_helpers.get_temp_export_dir(export_dir) builder = saved_model_builder._SavedModelBuilder(temp_export_dir) # Manually save variables to export them in an object-based checkpoint. This # skips the `builder.add_meta_graph_and_variables()` step, which saves a # named-based checkpoint. # TODO(b/113134168): Add fn to Builder to save with object-based saver. # TODO(b/113178242): This should only export the model json structure. Only # one save is needed once the weights can be copied from the model to clone. checkpoint_path = _export_model_json_and_variables(model, temp_export_dir) # Export each mode. Use ModeKeys enums defined for `Estimator` to ensure that # Keras models and `Estimator`s are exported with the same format. # Every time a mode is exported, the code checks to see if new variables have # been created (e.g. optimizer slot variables). If that is the case, the # checkpoint is re-saved to include the new variables. export_args = { 'builder': builder, 'model': model, 'custom_objects': custom_objects, 'checkpoint_path': checkpoint_path } has_saved_vars = False if model.optimizer: if isinstance(model.optimizer, optimizers.TFOptimizer): _export_mode(model_fn_lib.ModeKeys.TRAIN, has_saved_vars, **export_args) has_saved_vars = True _export_mode(model_fn_lib.ModeKeys.EVAL, has_saved_vars, **export_args) else: logging.warning( 'Model was compiled with an optimizer, but the optimizer is not from ' '`tf.train` (e.g. `tf.train.AdagradOptimizer`). Only the serving ' 'graph was exported. The train and evaluate graphs were not added to ' 'the SavedModel.') _export_mode(model_fn_lib.ModeKeys.PREDICT, has_saved_vars, **export_args) builder.save(as_text) gfile.Rename(temp_export_dir, export_dir) return export_dir