def test_revived_sequential(self):
model = keras.models.Sequential()
model.add(keras.layers.Dense(5, input_shape=(3,),
kernel_regularizer=regularizers.get('l2')))
model.add(keras.layers.Dense(2, kernel_regularizer=regularizers.get('l2')))
self.evaluate(variables.variables_initializer(model.variables))
saved_model_dir = self._save_model_dir()
model.save(saved_model_dir, save_format='tf')
loaded = keras_load.load(saved_model_dir)
self.assertLen(loaded.layers, 2)
self.assertLen(loaded.losses, 2)
loaded.pop()
self.assertLen(loaded.layers, 1)
self.assertLen(loaded.losses, 1)
loaded.add(keras.layers.Dense(2, kernel_regularizer=regularizers.get('l2')))
self.assertLen(loaded.layers, 2)
self.assertLen(loaded.losses, 2)
def test_model_save_and_load(self):
input_arr = np.random.random((1, 3)).astype(np.float32)
target_arr = np.random.random((1, 4)).astype(np.float32)
model = testing_utils.get_small_mlp(1, 4, input_dim=3)
model.layers[-1].activity_regularizer = regularizers.get('l2')
model.activity_regularizer = regularizers.get('l2')
model.compile(loss='mse', optimizer='rmsprop')
model.train_on_batch(input_arr, target_arr)
def callable_loss():
return math_ops.reduce_sum(model.weights[0])
model.add_loss(callable_loss)
saved_model_dir = self._save_model_dir()
tf_save.save(model, saved_model_dir)
loaded = keras_load.load(saved_model_dir)
self.evaluate(variables.variables_initializer(loaded.variables))
self.assertAllClose(self.evaluate(model.weights),
self.evaluate(loaded.weights))
input_arr = constant_op.constant(
np.random.random((1, 3)).astype(np.float32))
self.assertAllClose(self.evaluate(model(input_arr)),
self.evaluate(loaded(input_arr)))
# Validate losses. The order of conditional losses may change between the
# model and loaded model, so sort the losses first.
if context.executing_eagerly():
self.assertAllClose(sorted(self.evaluate(model.losses)),
sorted(self.evaluate(loaded.losses)))
else:
self.assertAllClose(self.evaluate(model.get_losses_for(None)),
self.evaluate(loaded.get_losses_for(None)))
self.assertAllClose(
sorted(self.evaluate(model.get_losses_for(input_arr))),
sorted(self.evaluate(loaded.get_losses_for(input_arr))))
def testSaveWithSignatures(self):
model = keras.models.Sequential()
model.add(
keras.layers.Dense(5,
input_shape=(3, ),
kernel_regularizer=regularizers.get('l2')))
model.add(keras.layers.Dropout(0.5))
model.add(
keras.layers.Dense(4, kernel_regularizer=regularizers.get('l2')))
input_arr = np.random.random((2, 3))
target_arr = np.random.random((2, 4))
model.compile(loss='mse', optimizer='rmsprop')
model.train_on_batch(input_arr, target_arr)
@def_function.function(
input_signature=[tensor_spec.TensorSpec((None, 3))])
def predict(inputs):
return {'predictions': model(inputs)}
feature_configs = {
'inputs':
parsing_ops.FixedLenFeature(shape=[2, 3], dtype=dtypes.float32)
}
@def_function.function(
input_signature=[tensor_spec.TensorSpec([None], dtypes.string)])
def parse_and_predict(examples):
features = parsing_ops.parse_single_example(
examples[0], feature_configs)
return {
'predictions': model(features['inputs']),
'layer_1_outputs': model.layers[0](features['inputs'])
}
saved_model_dir = self._save_model_dir()
model.save(saved_model_dir,
save_format='tf',
signatures={
'predict': predict,
'parse_and_predict': parse_and_predict
})
model.save('/tmp/saved',
save_format='tf',
signatures={
'predict': predict,
'parse_and_predict': parse_and_predict
})
loaded = keras_load.load(saved_model_dir)
self.assertAllClose(
model.predict(input_arr),
loaded.signatures['predict'](ops.convert_to_tensor_v2(
input_arr.astype('float32')))['predictions'])
feature = {
'inputs':
feature_pb2.Feature(float_list=feature_pb2.FloatList(
value=input_arr.astype('float32').flatten()))
}
example = example_pb2.Example(features=feature_pb2.Features(
feature=feature))
outputs = loaded.signatures['parse_and_predict'](
ops.convert_to_tensor_v2([example.SerializeToString()]))
self.assertAllClose(model.predict(input_arr), outputs['predictions'])
self.assertAllClose(model.layers[0](input_arr),
outputs['layer_1_outputs'])
def test_revive_network(self, model_cls): model = model_cls(8) model.save(self.path, include_optimizer=False, save_format='tf') revived = keras_load.load(self.path, compile=False) self._assert_revived_correctness(model, revived)
def test_must_restore_from_config_registration(self):
layer = GlobalLayerThatShouldFailIfNotAdded()
saved_model_dir = self._save_model_dir()
tf_save.save(layer, saved_model_dir)
_ = keras_load.load(saved_model_dir)
def _save_and_load(self, model): saved_model_dir = self._save_model_dir() tf_save.save(model, saved_model_dir) loaded = keras_load.load(saved_model_dir) return loaded
def test_revive(self):
input_shape = None
if testing_utils.get_model_type() == 'functional':
input_shape = (2, 3)
layer_with_config = CustomLayerWithConfig(1., 2)
layer_without_config = CustomLayerNoConfig(3., 4)
subclassed_with_config = SubclassedModelWithConfig(4., 6.)
subclassed_without_config = SubclassedModelNoConfig(7., 8.)
inputs = keras.Input((2, 3))
x = CustomLayerWithConfig(1., 2)(inputs)
x = CustomLayerNoConfig(3., 4)(x)
x = SubclassedModelWithConfig(4., 6.)(x)
x = SubclassedModelNoConfig(7., 8.)(x)
inner_model_functional = keras.Model(inputs, x)
inner_model_sequential = keras.Sequential([
CustomLayerWithConfig(1., 2),
CustomLayerNoConfig(3., 4),
SubclassedModelWithConfig(4., 6.),
SubclassedModelNoConfig(7., 8.)
])
class SubclassedModel(keras.Model):
def __init__(self):
super(SubclassedModel, self).__init__()
self.all_layers = [
CustomLayerWithConfig(1., 2),
CustomLayerNoConfig(3., 4),
SubclassedModelWithConfig(4., 6.),
SubclassedModelNoConfig(7., 8.)
]
def call(self, inputs):
x = inputs
for layer in self.all_layers:
x = layer(x)
return x
inner_model_subclassed = SubclassedModel()
layers = [
layer_with_config, layer_without_config, subclassed_with_config,
subclassed_without_config, inner_model_functional,
inner_model_sequential, inner_model_subclassed
]
model = testing_utils.get_model_from_layers(layers,
input_shape=input_shape)
# Run data through the Model to create save spec and weights.
model.predict(np.ones((10, 2, 3)), batch_size=10)
# Test that the correct checkpointed values are loaded, whether the layer is
# created from the config or SavedModel.
layer_with_config.c.assign(2 * layer_with_config.c)
layer_without_config.c.assign(3 * layer_without_config.c)
model.save(self.path, save_format='tf')
revived = keras_load.load(self.path)
self._assert_revived_correctness(model, revived)
def test_revive_subclassed_with_nested_model(self): model = SubclassedModelNoConfig(1., 2.) model._set_inputs(tensor_spec.TensorSpec((None, 2, 3))) model.save(self.path, save_format='tf') revived = keras_load.load(self.path) self._assert_revived_correctness(model, revived)
def test_revive(self):
input_shape = None
if testing_utils.get_model_type() == 'functional':
input_shape = (2, 3)
layer_with_config = CustomLayerWithConfig(1., 2)
layer_without_config = CustomLayerNoConfig(3., 4)
subclassed_with_config = SubclassedModelWithConfig(4., 6.)
subclassed_without_config = SubclassedModelNoConfig(7., 8.)
inputs = keras.Input((2, 3))
x = CustomLayerWithConfig(1., 2)(inputs)
x = CustomLayerNoConfig(3., 4)(x)
x = SubclassedModelWithConfig(4., 6.)(x)
x = SubclassedModelNoConfig(7., 8.)(x)
inner_model_functional = keras.Model(inputs, x)
inner_model_sequential = keras.Sequential(
[CustomLayerWithConfig(1., 2),
CustomLayerNoConfig(3., 4),
SubclassedModelWithConfig(4., 6.),
SubclassedModelNoConfig(7., 8.)])
class SubclassedModel(keras.Model):
def __init__(self):
super(SubclassedModel, self).__init__()
self.all_layers = [CustomLayerWithConfig(1., 2),
CustomLayerNoConfig(3., 4),
SubclassedModelWithConfig(4., 6.),
SubclassedModelNoConfig(7., 8.)]
def call(self, inputs):
x = inputs
for layer in self.all_layers:
x = layer(x)
return x
inner_model_subclassed = SubclassedModel()
layers = [layer_with_config,
layer_without_config,
subclassed_with_config,
subclassed_without_config,
inner_model_functional,
inner_model_sequential,
inner_model_subclassed]
model = testing_utils.get_model_from_layers(
layers, input_shape=input_shape)
# The inputs attribute must be defined in order to save the model.
if not model.inputs:
model._set_inputs(tensor_spec.TensorSpec((None, 2, 3)))
# Test that the correct checkpointed values are loaded, whether the layer is
# created from the config or SavedModel.
layer_with_config.c.assign(2 * layer_with_config.c)
layer_without_config.c.assign(3 * layer_without_config.c)
model.save(self.path, save_format='tf')
revived = keras_load.load(self.path)
self._assert_revived_correctness(model, revived)
