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'):
training_utils.trace_model_call(model)
fn = training_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'):
training_utils.trace_model_call(model)
model._set_inputs(inputs)
fn = training_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'):
training_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 = training_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 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 = training_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_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 = training_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 save_keras_model(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.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. 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 = export_helpers.get_timestamped_export_dir(saved_model_path)
if serving_only:
save_lib.save(model,
export_dir,
signatures=training_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
