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., s_bk_np + 10., 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(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.assertRaisesRegexp(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 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 context.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()
tf_save.save(model, saved_model_dir)
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 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_save.save(layer, saved_model_dir)
with generic_utils.CustomObjectScope(
{'LayerThatShouldFailIfNotAdded': LayerThatShouldFailIfNotAdded}):
_ = keras_load.load(saved_model_dir)
def test_custom_metric_wrapped_call(self):
class NegativeMean(keras.metrics.Mean):
@def_function.function(
input_signature=[tensor_spec.TensorSpec(None, dtypes.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 load_model(filepath, custom_objects=None, compile=True): # pylint: disable=redefined-builtin
"""Loads a model saved via `save_model`.
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`.
Arguments:
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.
Returns:
A Keras model instance. If an optimizer was found
as part of the saved model, the model is already
compiled. Otherwise, the model is uncompiled and
a warning will be displayed. When `compile` is set
to False, the compilation is omitted without any
warning.
Raises:
ImportError: if loading from an hdf5 file and h5py is not available.
IOError: In case of an invalid savefile.
"""
with generic_utils.CustomObjectScope(custom_objects or {}):
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)
if sys.version_info >= (3, 4) and isinstance(filepath, pathlib.Path):
filepath = str(filepath)
if isinstance(filepath, six.string_types):
loader_impl.parse_saved_model(filepath)
return saved_model_load.load(filepath, compile)
raise IOError(
'Unable to load model. Filepath is not an hdf5 file (or h5py is not '
'available) or SavedModel.')
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_load.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(constant_op.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(constant_op.constant([1.]))
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_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=dtypes.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, string_types):
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 = testing_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)
revived.test_on_batch(x, y_true))
# Compile with sparse categorical accuracy
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__':
ops.enable_eager_execution()
with generic_utils.CustomObjectScope({
'CustomLayerWithConfig': CustomLayerWithConfig,
'CustomNetworkWithConfig': CustomNetworkWithConfig,
'CustomNetworkWithConfigName': CustomNetworkWithConfigName,
'SubclassedModelWithConfig': SubclassedModelWithConfig
}):
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)
regularizer = (mp_test_util.IdentityRegularizer()
if use_regularizer else None)
with strategy_fn().scope():
# Pass loss_scale=None, as this test will fail if the DynamicLossScale
# skips applying gradients for a step
with policy.policy_scope(
policy.Policy(policy_name, loss_scale=None)):
layer = mp_test_util.MultiplyLayer(assert_type=dtypes.float16,
use_operator=use_operator,
regularizer=regularizer,
input_shape=(1, ))
if use_input_spec:
layer.input_spec = input_spec.InputSpec(shape=(2, 1))
model = testing_utils.get_model_from_layers(
[layer], input_shape=(1, ), input_dtype=dtypes.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 math_ops.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)
model.compile(opt,
loss=loss_fn,
run_eagerly=testing_utils.should_run_eagerly())
x = np.ones((2, 1))
y = np.ones((2, 1))
dataset = dataset_ops.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:
# Regularizer adds another 2 ** -14 to the gradient.
expected -= 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)
def load(
self,
*,
timestamp: Optional[Timestamp] = None,
compile_model: bool = False,
custom_objects: Optional[Mapping[str, Any]] = None,
input_shape: Optional[Tuple[int, ...]] = None,
) -> tf.keras.Model:
"""
Load a Tensorflow model from a TileDB array.
:param timestamp: Range of timestamps to load fragments of the array which live
in the specified time range.
:param compile_model: Whether to compile the model after loading or not.
:param custom_objects: Mapping of names to custom classes or functions to be
considered during deserialization.
:param input_shape: The shape that the custom model expects as input
:return: Tensorflow model.
"""
# TODO: Change timestamp when issue in core is resolved
with tiledb.open(self.uri, ctx=self.ctx,
timestamp=timestamp) as model_array:
model_array_results = model_array[:]
model_config = json.loads(model_array.meta["model_config"])
model_class = model_config["class_name"]
if model_class != "Sequential" and model_class != "Functional":
with generic_utils.SharedObjectLoadingScope():
with generic_utils.CustomObjectScope(custom_objects or {}):
if hasattr(model_config, "decode"):
model_config = model_config.decode("utf-8")
model = model_config_lib.model_from_config(
model_config, custom_objects=custom_objects)
if not model.built:
model.build(input_shape)
# Load weights for layers
self._load_custom_subclassed_model(model, model_array)
else:
cls = (tf.keras.Sequential
if model_class == "Sequential" else tf.keras.Model)
model = cls.from_config(model_config["config"])
model_weights = pickle.loads(
model_array_results["model_weights"].item(0))
model.set_weights(model_weights)
if compile_model:
optimizer_weights = pickle.loads(
model_array_results["optimizer_weights"].item(0))
training_config = json.loads(
model_array.meta["training_config"])
# Compile model.
model.compile(**saving_utils.compile_args_from_training_config(
training_config, custom_objects))
saving_utils.try_build_compiled_arguments(model)
# Set optimizer weights.
if optimizer_weights:
try:
model.optimizer._create_all_weights(
model.trainable_variables)
except (NotImplementedError, AttributeError):
logging.warning(
"Error when creating the weights of optimizer {}, making it "
"impossible to restore the saved optimizer state. As a result, "
"your model is starting with a freshly initialized optimizer."
)
try:
model.optimizer.set_weights(optimizer_weights)
except ValueError:
logging.warning("Error in loading the saved optimizer "
"state. As a result, your model is "
"starting with a freshly initialized "
"optimizer.")
return model
