def test_Bidirectional_with_constants_layer_passing_initial_state(self):
with self.cached_session():
# Test basic case.
x = keras.Input((5, 5))
c = keras.Input((3, ))
s_for = keras.Input((32, ))
s_bac = keras.Input((32, ))
cell = _RNNCellWithConstants(32, 3)
custom_objects = {"_RNNCellWithConstants": _RNNCellWithConstants}
with generic_utils.CustomObjectScope(custom_objects):
layer = keras.layers.Bidirectional(keras.layers.RNN(cell))
y = layer(x, initial_state=[s_for, s_bac], constants=c)
model = keras.Model([x, s_for, s_bac, c], y)
model.compile(optimizer="rmsprop", loss="mse")
model.train_on_batch(
[
np.zeros((6, 5, 5)),
np.zeros((6, 32)),
np.zeros((6, 32)),
np.zeros((6, 3)),
],
np.zeros((6, 64)),
)
# Test basic case serialization.
x_np = np.random.random((6, 5, 5))
s_fw_np = np.random.random((6, 32))
s_bk_np = np.random.random((6, 32))
c_np = np.random.random((6, 3))
y_np = model.predict([x_np, s_fw_np, s_bk_np, c_np])
weights = model.get_weights()
config = layer.get_config()
with generic_utils.CustomObjectScope(custom_objects):
layer = keras.layers.Bidirectional.from_config(
copy.deepcopy(config))
y = layer(x, initial_state=[s_for, s_bac], constants=c)
model = keras.Model([x, s_for, s_bac, c], y)
model.set_weights(weights)
y_np_2 = model.predict([x_np, s_fw_np, s_bk_np, c_np])
self.assertAllClose(y_np, y_np_2, atol=1e-4)
# Verify that state is used
y_np_2_different_s = model.predict(
[x_np, s_fw_np + 10.0, s_bk_np + 10.0, c_np])
assert np.mean(y_np - y_np_2_different_s) != 0
# Test flat list inputs
with generic_utils.CustomObjectScope(custom_objects):
layer = keras.layers.Bidirectional.from_config(
copy.deepcopy(config))
y = layer([x, s_for, s_bac, c])
model = keras.Model([x, s_for, s_bac, c], y)
model.set_weights(weights)
y_np_3 = model.predict([x_np, s_fw_np, s_bk_np, c_np])
self.assertAllClose(y_np, y_np_3, atol=1e-4)
def test_custom_metric_model(self):
# TODO(b/134519980): Issue with `model.fit` if the model call function uses
# a `tf.function` in graph mode.
if not tf.executing_eagerly():
return
x = np.random.random((1, 3))
y = np.random.random((1, 4))
class CustomMetric(keras.metrics.MeanSquaredError):
pass
def zero_metric(y_true, y_pred):
del y_true, y_pred
return 0
model = testing_utils.get_small_mlp(1, 4, input_dim=3)
model.compile(loss='mse', optimizer='SGD',
metrics=[CustomMetric(), zero_metric])
model.fit(x, y)
saved_model_dir = self._save_model_dir()
model.save(saved_model_dir, save_format='tf')
with self.assertRaisesRegex(ValueError, 'custom_objects'):
keras_load.load(saved_model_dir)
with generic_utils.CustomObjectScope(
{'CustomMetric': CustomMetric, 'zero_metric': zero_metric}):
loaded = keras_load.load(saved_model_dir)
self.evaluate([v.initializer for v in loaded.variables])
loaded.fit(x, y)
def test_custom_metric(self, base_cls, num_tensor_args, requires_build):
class CustomMetric(base_cls):
def update_state(self, *args): # pylint: disable=useless-super-delegation
# Sometimes built-in metrics return an op in update_state. Custom
# metrics don't support returning ops, so wrap the update_state method
# while returning nothing.
super(CustomMetric, self).update_state(*args)
with self.cached_session():
metric = CustomMetric()
save_dir = self._save_model_dir('first_save')
if requires_build:
metric(*self.generate_inputs(num_tensor_args)) # pylint: disable=not-callable
self.evaluate([v.initializer for v in metric.variables])
with self.assertRaisesRegex(ValueError,
'Unable to restore custom object'):
self._test_metric_save_and_load(metric, save_dir, num_tensor_args)
with generic_utils.CustomObjectScope({'CustomMetric': CustomMetric}):
loaded = self._test_metric_save_and_load(
metric,
save_dir,
num_tensor_args,
test_sample_weight=False)
self._test_metric_save_and_load(
loaded,
self._save_model_dir('second_save'),
num_tensor_args,
test_sample_weight=False)
def load_model(filepath, custom_objects=None, compile=True, options=None): # pylint: disable=redefined-builtin
"""Loads a model saved via `model.save()`.
Usage:
>>> model = tf.keras.Sequential([
... tf.keras.layers.Dense(5, input_shape=(3,)),
... tf.keras.layers.Softmax()])
>>> model.save('/tmp/model')
>>> loaded_model = tf.keras.models.load_model('/tmp/model')
>>> x = tf.random.uniform((10, 3))
>>> assert np.allclose(model.predict(x), loaded_model.predict(x))
Note that the model weights may have different scoped names after being
loaded. Scoped names include the model/layer names, such as
`"dense_1/kernel:0"`. It is recommended that you use the layer properties to
access specific variables, e.g. `model.get_layer("dense_1").kernel`.
Args:
filepath: One of the following:
- String or `pathlib.Path` object, path to the saved model
- `h5py.File` object from which to load the model
custom_objects: Optional dictionary mapping names
(strings) to custom classes or functions to be
considered during deserialization.
compile: Boolean, whether to compile the model
after loading.
options: Optional `tf.saved_model.LoadOptions` object that specifies
options for loading from SavedModel.
Returns:
A Keras model instance. If the original model was compiled, and saved with
the optimizer, then the returned model will be compiled. Otherwise, the
model will be left uncompiled. In the case that an uncompiled model is
returned, a warning is displayed if the `compile` argument is set to
`True`.
Raises:
ImportError: if loading from an hdf5 file and h5py is not available.
IOError: In case of an invalid savefile.
"""
with generic_utils.SharedObjectLoadingScope():
with generic_utils.CustomObjectScope(custom_objects or {}):
with load_context.load_context(options):
if (h5py is not None and (isinstance(filepath, h5py.File)
or h5py.is_hdf5(filepath))):
return hdf5_format.load_model_from_hdf5(
filepath, custom_objects, compile)
filepath = path_to_string(filepath)
if isinstance(filepath, six.string_types):
loader_impl.parse_saved_model(filepath)
return saved_model_load.load(filepath, compile, options)
raise IOError(
'Unable to load model. Filepath is not an hdf5 file (or h5py is not '
'available) or SavedModel.')
def test_must_restore_from_config_custom_object_scope(self):
class LayerThatShouldFailIfNotAdded(keras.layers.Layer):
_must_restore_from_config = True
layer = LayerThatShouldFailIfNotAdded()
saved_model_dir = self._save_model_dir()
tf.saved_model.save(layer, saved_model_dir)
with generic_utils.CustomObjectScope(
{'LayerThatShouldFailIfNotAdded': LayerThatShouldFailIfNotAdded}):
_ = keras_load.load(saved_model_dir)
def test_Bidirectional_with_constants(self):
with self.cached_session():
# Test basic case.
x = keras.Input((5, 5))
c = keras.Input((3, ))
cell = _RNNCellWithConstants(32, 3)
custom_objects = {"_RNNCellWithConstants": _RNNCellWithConstants}
with generic_utils.CustomObjectScope(custom_objects):
layer = keras.layers.Bidirectional(keras.layers.RNN(cell))
y = layer(x, constants=c)
model = keras.Model([x, c], y)
model.compile(optimizer="rmsprop", loss="mse")
model.train_on_batch([np.zeros(
(6, 5, 5)), np.zeros((6, 3))], np.zeros((6, 64)))
# Test basic case serialization.
x_np = np.random.random((6, 5, 5))
c_np = np.random.random((6, 3))
y_np = model.predict([x_np, c_np])
weights = model.get_weights()
config = layer.get_config()
with generic_utils.CustomObjectScope(custom_objects):
layer = keras.layers.Bidirectional.from_config(
copy.deepcopy(config))
y = layer(x, constants=c)
model = keras.Model([x, c], y)
model.set_weights(weights)
y_np_2 = model.predict([x_np, c_np])
self.assertAllClose(y_np, y_np_2, atol=1e-4)
# Test flat list inputs
with generic_utils.CustomObjectScope(custom_objects):
layer = keras.layers.Bidirectional.from_config(
copy.deepcopy(config))
y = layer([x, c])
model = keras.Model([x, c], y)
model.set_weights(weights)
y_np_3 = model.predict([x_np, c_np])
self.assertAllClose(y_np, y_np_3, atol=1e-4)
def test_must_restore_from_config_custom_object_scope(self):
class LayerThatShouldFailIfNotAdded(keras.layers.Layer):
_must_restore_from_config = True
layer = LayerThatShouldFailIfNotAdded()
model = testing_utils.get_model_from_layers(
[layer], input_shape=[3], model_type='functional')
saved_model_dir = self._save_model_dir()
model.save(saved_model_dir, save_format='tf')
with generic_utils.CustomObjectScope(
{'LayerThatShouldFailIfNotAdded': LayerThatShouldFailIfNotAdded}):
_ = keras_load.load(saved_model_dir)
def test_custom_metric_wrapped_call(self):
class NegativeMean(keras.metrics.Mean):
@tf.function(input_signature=[tf.TensorSpec(None, tf.float32)])
def update_state(self, value):
super(NegativeMean, self).update_state(-value)
metric = NegativeMean()
self.evaluate([v.initializer for v in metric.variables])
with generic_utils.CustomObjectScope({'NegativeMean': NegativeMean}):
self._test_metric_save_and_load(metric,
self._save_model_dir(),
1,
test_sample_weight=False)
def compile_args_from_training_config(training_config, custom_objects=None):
"""Return model.compile arguments from training config."""
if custom_objects is None:
custom_objects = {}
with generic_utils.CustomObjectScope(custom_objects):
optimizer_config = training_config["optimizer_config"]
optimizer = optimizers.deserialize(optimizer_config)
# Recover losses.
loss = None
loss_config = training_config.get("loss", None)
if loss_config is not None:
loss = _deserialize_nested_config(losses.deserialize, loss_config)
# Recover metrics.
metrics = None
metrics_config = training_config.get("metrics", None)
if metrics_config is not None:
metrics = _deserialize_nested_config(
_deserialize_metric, metrics_config
)
# Recover weighted metrics.
weighted_metrics = None
weighted_metrics_config = training_config.get("weighted_metrics", None)
if weighted_metrics_config is not None:
weighted_metrics = _deserialize_nested_config(
_deserialize_metric, weighted_metrics_config
)
sample_weight_mode = (
training_config["sample_weight_mode"]
if hasattr(training_config, "sample_weight_mode")
else None
)
loss_weights = training_config["loss_weights"]
return dict(
optimizer=optimizer,
loss=loss,
metrics=metrics,
weighted_metrics=weighted_metrics,
loss_weights=loss_weights,
sample_weight_mode=sample_weight_mode,
)
def test_save_without_tracing(self):
class DoNotTrace(keras.layers.Layer):
def __init__(self):
super(DoNotTrace, self).__init__()
self.input_spec = keras.layers.InputSpec(shape=[None])
self.built = True
def call(self, inputs):
raise ValueError('I said do not trace')
def get_config(self):
return {}
@property
def _use_input_spec_as_call_signature(self):
return True
root = keras.models.Sequential()
root.add(keras.layers.Input(shape=(3,)))
root.attached_layer = DoNotTrace()
saved_model_dir = self._save_model_dir()
# With the default settings, the call function is traced.
with self.assertRaisesRegex(ValueError, 'do not trace'):
root.save(saved_model_dir, save_format='tf')
# When saving the config only, the layer call function should not be not
# traced.
root.save(saved_model_dir, save_format='tf', save_traces=False)
loaded = tf.saved_model.load(saved_model_dir)
self.assertTrue(hasattr(loaded, 'attached_layer'))
# This should raise an error when loaded without the custom object
loaded = keras_load.load(saved_model_dir)
with self.assertRaisesRegex(ValueError, 'Cannot call custom layer'):
loaded.attached_layer(tf.constant([1.]))
# Try loading with the custom objects
with generic_utils.CustomObjectScope({'DoNotTrace': DoNotTrace}):
loaded = keras_load.load(saved_model_dir)
with self.assertRaisesRegex(ValueError, 'I said do not trace'):
loaded.attached_layer(tf.constant([1.]))
def test_shared_objects(self):
class OuterLayer(keras.layers.Layer):
def __init__(self, inner_layer):
super(OuterLayer, self).__init__()
self.inner_layer = inner_layer
def call(self, inputs):
return self.inner_layer(inputs)
def get_config(self):
return {
'inner_layer': generic_utils.serialize_keras_object(
self.inner_layer)
}
@classmethod
def from_config(cls, config):
return cls(generic_utils.deserialize_keras_object(
config['inner_layer']))
class InnerLayer(keras.layers.Layer):
def __init__(self):
super(InnerLayer, self).__init__()
self.v = self.add_weight(name='v', shape=[], dtype=tf.float32)
def call(self, inputs):
return self.v + inputs
@classmethod
def from_config(cls, config):
return cls()
# Create a model with 2 output layers that share the same inner layer.
inner_layer = InnerLayer()
outer_layer_1 = OuterLayer(inner_layer)
outer_layer_2 = OuterLayer(inner_layer)
input_ = keras.Input(shape=(1,))
model = keras.Model(
inputs=input_, outputs=[outer_layer_1(input_), outer_layer_2(input_)])
# Changes to the shared layer should affect both outputs.
model.layers[1].inner_layer.v.assign(5)
self.assertAllEqual(model(1), [6.0, 6.0])
model.layers[1].inner_layer.v.assign(3)
self.assertAllEqual(model(1), [4.0, 4.0])
# After loading, changes to the shared layer should still affect both
# outputs.
def _do_assertions(loaded):
loaded.layers[1].inner_layer.v.assign(5)
self.assertAllEqual(loaded(1), [6.0, 6.0])
loaded.layers[1].inner_layer.v.assign(3)
self.assertAllEqual(loaded(1), [4.0, 4.0])
loaded.layers[2].inner_layer.v.assign(5)
self.assertAllEqual(loaded(1), [6.0, 6.0])
loaded.layers[2].inner_layer.v.assign(3)
self.assertAllEqual(loaded(1), [4.0, 4.0])
# We'd like to make sure we only attach shared object IDs when strictly
# necessary, so we'll recursively traverse the generated config to count
# whether we have the exact number we expect.
def _get_all_keys_recursive(dict_or_iterable):
if isinstance(dict_or_iterable, dict):
for key in dict_or_iterable.keys():
yield key
for key in _get_all_keys_recursive(dict_or_iterable.values()):
yield key
elif isinstance(dict_or_iterable, str):
return
else:
try:
for item in dict_or_iterable:
for key in _get_all_keys_recursive(item):
yield key
# Not an iterable or dictionary
except TypeError:
return
with generic_utils.CustomObjectScope({
'OuterLayer': OuterLayer, 'InnerLayer': InnerLayer}):
# Test saving and loading to disk
save_format = test_utils.get_save_format()
saved_model_dir = self._save_model_dir()
keras.models.save_model(model, saved_model_dir, save_format=save_format)
loaded = keras.models.load_model(saved_model_dir)
_do_assertions(loaded)
# Test recreating directly from config
config = model.get_config()
key_count = collections.Counter(_get_all_keys_recursive(config))
self.assertEqual(key_count[generic_utils.SHARED_OBJECT_KEY], 2)
loaded = keras.Model.from_config(config)
_do_assertions(loaded)
model.compile('rmsprop', 'mse', 'acc')
y_true = np.random.randint(0, 3, (5, 1)).astype(np.float32)
model.train_on_batch(x, y_true)
model.save(self.path, include_optimizer=True, save_format='tf')
revived = keras_load.load(self.path, compile=True)
self.assertAllClose(model.test_on_batch(x, y_true),
revived.test_on_batch(x, y_true))
def test_revived_model_has_save_spec(self):
model = SubclassedModelWithConfig(2, 3)
model.predict(np.random.random((5, 10)).astype(np.float32))
model.save(self.path, save_format='tf')
revived = keras_load.load(self.path, compile=True)
self.assertAllEqual(model._get_save_spec(dynamic_batch=False),
revived._get_save_spec(dynamic_batch=False))
if __name__ == '__main__':
tf.compat.v1.enable_eager_execution()
with generic_utils.CustomObjectScope({
'CustomLayerWithConfig':
CustomLayerWithConfig,
'CustomNetworkWithConfig':
CustomNetworkWithConfig,
'CustomNetworkWithConfigName':
CustomNetworkWithConfigName,
'SubclassedModelWithConfig':
SubclassedModelWithConfig
}):
tf.test.main()
def test_model(
self,
strategy_fn,
use_operator=False,
use_regularizer=False,
policy_name="mixed_float16",
get_config=False,
save_format=None,
use_input_spec=False,
):
self._skip_if_strategy_unsupported(strategy_fn)
self._skip_if_save_format_unsupported(save_format)
if use_regularizer:
weight_regularizer = mp_test_util.IdentityRegularizer()
activity_regularizer = mp_test_util.ReduceSumRegularizer()
else:
weight_regularizer = activity_regularizer = None
with strategy_fn().scope():
with policy.policy_scope(policy_name):
layer = mp_test_util.MultiplyLayer(
assert_type=tf.float16,
use_operator=use_operator,
regularizer=weight_regularizer,
activity_regularizer=activity_regularizer,
input_shape=(1,),
)
if use_input_spec:
layer.input_spec = input_spec.InputSpec(shape=(None, 1))
model = test_utils.get_model_from_layers(
[layer], input_shape=(1,), input_dtype=tf.float16
)
if get_config:
config = model.get_config()
model = model.__class__.from_config(
config,
custom_objects={
"MultiplyLayer": mp_test_util.MultiplyLayer
},
)
(layer,) = (
layer
for layer in model.layers
if isinstance(layer, mp_test_util.MultiplyLayer)
)
def loss_fn(y_true, y_pred):
del y_true
return tf.reduce_mean(y_pred)
# Learning rate is small enough that if applied to a float16 variable,
# the variable will not change. So this tests the learning rate not
# applied to a float16 value, but instead the float32 variable.
opt = gradient_descent.SGD(2**-14)
# Use a fixed loss scale, as this test will fail if gradients are
# skipped for a step due to dynamic loss scaling.
opt = loss_scale_optimizer.LossScaleOptimizer(
opt, dynamic=False, initial_scale=8
)
model.compile(
opt,
loss=loss_fn,
run_eagerly=test_utils.should_run_eagerly(),
)
x = np.ones((2, 1))
y = np.ones((2, 1))
dataset = tf.data.Dataset.from_tensor_slices((x, y)).batch(2)
model.fit(dataset)
# Variable starts at 1, and should have gradient of 2 ** -14 subtracted
# from it.
expected = 1 - 2**-14
if use_regularizer:
# Weight and activity regularizer each add another 2 ** -14 to the
# gradient.
expected -= 2 * 2**-14
self.assertEqual(backend.eval(layer.v), expected)
if save_format:
with generic_utils.CustomObjectScope(
{
"MultiplyLayer": mp_test_util.MultiplyLayer,
"loss_fn": loss_fn,
}
):
self._test_saving(model, dataset, save_format, use_regularizer)
pass
model = LinearModel(2, 3)
model(np.random.random((5, 10)).astype(np.float32))
model.save(self.path, save_format='tf')
with self.assertRaisesRegex(
RuntimeError,
'Unable to restore object of class \'LinearModel\''):
keras_load.load(self.path, compile=True)
def test_load_model_with_name_conflict_registered_works(self):
model = WideDeepModel(2, 3)
model(np.random.random((5, 10)).astype(np.float32))
model.save(self.path, save_format='tf')
keras_load.load(self.path, compile=True)
if __name__ == '__main__':
tf.compat.v1.enable_eager_execution()
with generic_utils.CustomObjectScope({
'CustomLayerWithConfig': CustomLayerWithConfig,
'CustomNetworkWithConfig': CustomNetworkWithConfig,
'CustomNetworkWithConfigName': CustomNetworkWithConfigName,
'SubclassedModelWithConfig': SubclassedModelWithConfig,
'FunctionalSubclassModel': FunctionalSubclassModel,
'FunctionalSubclassModelWrongConfig':
FunctionalSubclassModelWrongConfig,
'WideDeepModel': WideDeepModel
}):
tf.test.main()
pass
model = LinearModel(2, 3)
model(np.random.random((5, 10)).astype(np.float32))
model.save(self.path, save_format="tf")
with self.assertRaisesRegex(
RuntimeError,
"Unable to restore object of class 'LinearModel'"):
keras_load.load(self.path, compile=True)
def test_load_model_with_name_conflict_registered_works(self):
model = WideDeepModel(2, 3)
model(np.random.random((5, 10)).astype(np.float32))
model.save(self.path, save_format="tf")
keras_load.load(self.path, compile=True)
if __name__ == "__main__":
tf.compat.v1.enable_eager_execution()
with generic_utils.CustomObjectScope({
"CustomLayerWithConfig": CustomLayerWithConfig,
"CustomNetworkWithConfig": CustomNetworkWithConfig,
"CustomNetworkWithConfigName": CustomNetworkWithConfigName,
"SubclassedModelWithConfig": SubclassedModelWithConfig,
"FunctionalSubclassModel": FunctionalSubclassModel,
"FunctionalSubclassModelWrongConfig":
FunctionalSubclassModelWrongConfig,
"WideDeepModel": WideDeepModel,
}):
tf.test.main()
