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
