def export_saved_model(estimator, export_dir_base, checkpoint_path, serving_input_receiver_fn, as_text=False): with context.graph_mode(): export_dir = export_helpers.get_timestamped_export_dir(export_dir_base) temp_export_dir = export_helpers.get_temp_export_dir(export_dir) builder = saved_model_builder.SavedModelBuilder(temp_export_dir) save_variables = True _add_meta_graph_for_mode(estimator, builder, serving_input_receiver_fn, checkpoint_path, save_variables) save_variables = False builder.save(as_text) if save_variables: raise ValueError('No valid modes for exporting found.') gfile.Rename(temp_export_dir, export_dir) return export_dir
def export_savedmodel( self, export_dir_base, serving_input_receiver_fn, assets_extra=None, as_text=False, checkpoint_path=None): """Exports inference graph as a SavedModel into given dir. This method builds a new graph by first calling the serving_input_receiver_fn to obtain feature `Tensor`s, and then calling this `Estimator`'s model_fn to generate the model graph based on those features. It restores the given checkpoint (or, lacking that, the most recent checkpoint) into this graph in a fresh session. Finally it creates a timestamped export directory below the given export_dir_base, and writes a `SavedModel` into it containing a single `MetaGraphDef` saved from this session. The exported `MetaGraphDef` will provide one `SignatureDef` for each element of the export_outputs dict returned from the model_fn, named using the same keys. One of these keys is always signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY, indicating which signature will be served when a serving request does not specify one. For each signature, the outputs are provided by the corresponding `ExportOutput`s, and the inputs are always the input receivers provided by the serving_input_receiver_fn. Extra assets may be written into the SavedModel via the extra_assets argument. This should be a dict, where each key gives a destination path (including the filename) relative to the assets.extra directory. The corresponding value gives the full path of the source file to be copied. For example, the simple case of copying a single file without renaming it is specified as `{'my_asset_file.txt': '/path/to/my_asset_file.txt'}`. Args: export_dir_base: A string containing a directory in which to create timestamped subdirectories containing exported SavedModels. serving_input_receiver_fn: A function that takes no argument and returns a `ServingInputReceiver`. assets_extra: A dict specifying how to populate the assets.extra directory within the exported SavedModel, or `None` if no extra assets are needed. as_text: whether to write the SavedModel proto in text format. checkpoint_path: The checkpoint path to export. If `None` (the default), the most recent checkpoint found within the model directory is chosen. Returns: The string path to the exported directory. Raises: ValueError: if no serving_input_receiver_fn is provided, no export_outputs are provided, or no checkpoint can be found. """ if serving_input_receiver_fn is None: raise ValueError('serving_input_receiver_fn must be defined.') with ops.Graph().as_default() as g: self._create_and_assert_global_step(g) random_seed.set_random_seed(self._config.tf_random_seed) serving_input_receiver = serving_input_receiver_fn() # Call the model_fn and collect the export_outputs. estimator_spec = self._call_model_fn( features=serving_input_receiver.features, labels=None, mode=model_fn_lib.ModeKeys.PREDICT, config=self.config) # Build the SignatureDefs from receivers and all outputs signature_def_map = build_all_signature_defs( serving_input_receiver.receiver_tensors, estimator_spec.export_outputs, serving_input_receiver.receiver_tensors_alternatives) if not checkpoint_path: # Locate the latest checkpoint checkpoint_path = saver.latest_checkpoint(self._model_dir) if not checkpoint_path: raise ValueError("Couldn't find trained model at %s." % self._model_dir) export_dir = get_timestamped_export_dir(export_dir_base) temp_export_dir = get_temp_export_dir(export_dir) # TODO(soergel): Consider whether MonitoredSession makes sense here with tf_session.Session() as session: saver_for_restore = estimator_spec.scaffold.saver or saver.Saver( sharded=True) saver_for_restore.restore(session, checkpoint_path) # TODO(b/36111876): replace legacy_init_op with main_op mechanism # pylint: disable=protected-access local_init_op = ( estimator_spec.scaffold.local_init_op or monitored_session.Scaffold._default_local_init_op()) # pylint: enable=protected-access # Perform the export builder = saved_model_builder.SavedModelBuilder(temp_export_dir) builder.add_meta_graph_and_variables( session, [tag_constants.SERVING], signature_def_map=signature_def_map, assets_collection=ops.get_collection( ops.GraphKeys.ASSET_FILEPATHS), legacy_init_op=local_init_op) builder.save(as_text) # Add the extra assets if assets_extra: assets_extra_path = os.path.join(compat.as_bytes(temp_export_dir), compat.as_bytes('assets.extra')) for dest_relative, source in assets_extra.items(): dest_absolute = os.path.join(compat.as_bytes(assets_extra_path), compat.as_bytes(dest_relative)) dest_path = os.path.dirname(dest_absolute) gfile.MakeDirs(dest_path) gfile.Copy(source, dest_absolute) gfile.Rename(temp_export_dir, export_dir) return export_dir
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. 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
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