From 1b7c53d0505b9fca65a98a7400a991cea982d1b1 Mon Sep 17 00:00:00 2001 From: Neel Kovelamudi Date: Fri, 28 Apr 2023 10:25:21 -0700 Subject: [PATCH] Adds Keras v3 saving testing coverage to Keras layers tests. PiperOrigin-RevId: 527921888 --- .../attention/multi_head_attention_test.py | 21 +- .../spectral_normalization_test.py | 25 ++- .../preprocessing/hashed_crossing_test.py | 44 +++- keras/layers/preprocessing/hashing_test.py | 24 +++ .../layers/preprocessing/index_lookup_test.py | 190 ++++++++++++------ .../preprocessing/integer_lookup_test.py | 67 ++++-- .../preprocessing/normalization_test.py | 28 ++- keras/layers/regularization/dropout_test.py | 103 ++++++---- 8 files changed, 359 insertions(+), 143 deletions(-) diff --git a/keras/layers/attention/multi_head_attention_test.py b/keras/layers/attention/multi_head_attention_test.py index 96b939ccd24..e9508cf86f4 100644 --- a/keras/layers/attention/multi_head_attention_test.py +++ b/keras/layers/attention/multi_head_attention_test.py @@ -19,6 +19,7 @@ from absl.testing import parameterized import keras +from keras.saving import object_registration from keras.testing_infra import test_combinations from keras.testing_infra import test_utils @@ -515,6 +516,7 @@ def test_initializer(self): self.assertEqual(output.shape.as_list(), [None, 40, 80]) +@object_registration.register_keras_serializable() class TestModel(keras.Model): def __init__(self): super().__init__() @@ -540,12 +542,19 @@ def call(self, x, training=False): @test_combinations.run_all_keras_modes(always_skip_v1=True) class KerasModelSavingTest(test_combinations.TestCase): - def test_keras_saving_subclass(self): + @parameterized.parameters("tf", "keras_v3") + def test_keras_saving_subclass(self, save_format): model = TestModel() query = keras.Input(shape=(40, 80)) _ = model(query) model_path = self.get_temp_dir() + "/tmp_model" - keras.models.save_model(model, model_path, save_format="tf") + if save_format == "keras_v3": + if not tf.__internal__.tf2.enabled(): + self.skipTest( + "TF2 must be enabled to use the new `.keras` saving." + ) + model_path += ".keras" + keras.models.save_model(model, model_path, save_format=save_format) reloaded_model = keras.models.load_model(model_path) self.assertEqual( len(model.trainable_variables), @@ -556,7 +565,7 @@ def test_keras_saving_subclass(self): ): self.assertAllEqual(src_v, loaded_v) - @parameterized.parameters("h5", "tf") + @parameterized.parameters("h5", "tf", "keras_v3") def test_keras_saving_functional(self, save_format): model = TestModel() query = keras.Input(shape=(40, 80)) @@ -565,6 +574,12 @@ def test_keras_saving_functional(self, save_format): )(query, query) model = keras.Model(inputs=query, outputs=output) model_path = self.get_temp_dir() + "/tmp_model" + if save_format == "keras_v3": + if not tf.__internal__.tf2.enabled(): + self.skipTest( + "TF2 must be enabled to use the new `.keras` saving." + ) + model_path += ".keras" keras.models.save_model(model, model_path, save_format=save_format) reloaded_model = keras.models.load_model(model_path) self.assertEqual( diff --git a/keras/layers/normalization/spectral_normalization_test.py b/keras/layers/normalization/spectral_normalization_test.py index 8d673879cd6..555850291af 100644 --- a/keras/layers/normalization/spectral_normalization_test.py +++ b/keras/layers/normalization/spectral_normalization_test.py @@ -51,12 +51,27 @@ def test_save_load_model(self): # initialize model model.predict(tf.random.uniform((2, 1))) - model.save("test.h5") - new_model = keras.models.load_model("test.h5") + with self.subTest("h5"): + model.save("test.h5") + new_model = keras.models.load_model("test.h5") - self.assertEqual( - model.layers[0].get_config(), new_model.layers[0].get_config() - ) + self.assertEqual( + model.layers[0].get_config(), new_model.layers[0].get_config() + ) + with self.subTest("savedmodel"): + model.save("test") + new_model = keras.models.load_model("test") + + self.assertEqual( + model.layers[0].get_config(), new_model.layers[0].get_config() + ) + with self.subTest("keras_v3"): + model.save("test.keras") + new_model = keras.models.load_model("test.keras") + + self.assertEqual( + model.layers[0].get_config(), new_model.layers[0].get_config() + ) @test_combinations.run_all_keras_modes def test_normalization(self): diff --git a/keras/layers/preprocessing/hashed_crossing_test.py b/keras/layers/preprocessing/hashed_crossing_test.py index 948dda50c32..6fa5163fb78 100644 --- a/keras/layers/preprocessing/hashed_crossing_test.py +++ b/keras/layers/preprocessing/hashed_crossing_test.py @@ -154,7 +154,7 @@ def test_from_config(self): tf.sparse.to_dense(original_outputs), ) - def test_saved_model_keras(self): + def test_saving_keras(self): string_in = keras.Input(shape=(1,), dtype=tf.string) int_in = keras.Input(shape=(1,), dtype=tf.int64) out = hashed_crossing.HashedCrossing(num_bins=10)((string_in, int_in)) @@ -167,17 +167,39 @@ def test_saved_model_keras(self): output_data = model((string_data, int_data)) self.assertAllClose(output_data, expected_output) - # Save the model to disk. - output_path = os.path.join(self.get_temp_dir(), "saved_model") - model.save(output_path, save_format="tf") - loaded_model = keras.models.load_model( - output_path, - custom_objects={"HashedCrossing": hashed_crossing.HashedCrossing}, - ) + with self.subTest("savedmodel"): + # Save the model to disk. + output_path = os.path.join(self.get_temp_dir(), "saved_model") + model.save(output_path, save_format="tf") + loaded_model = keras.models.load_model( + output_path, + custom_objects={ + "HashedCrossing": hashed_crossing.HashedCrossing + }, + ) + + # Validate correctness of the new model. + new_output_data = loaded_model((string_data, int_data)) + self.assertAllClose(new_output_data, expected_output) + + with self.subTest("keras_v3"): + if not tf.__internal__.tf2.enabled(): + self.skipTest( + "TF2 must be enabled to use the new `.keras` saving." + ) + # Save the model to disk. + output_path = os.path.join(self.get_temp_dir(), "model.keras") + model.save(output_path, save_format="keras_v3") + loaded_model = keras.models.load_model( + output_path, + custom_objects={ + "HashedCrossing": hashed_crossing.HashedCrossing + }, + ) - # Validate correctness of the new model. - new_output_data = loaded_model((string_data, int_data)) - self.assertAllClose(new_output_data, expected_output) + # Validate correctness of the new model. + new_output_data = loaded_model((string_data, int_data)) + self.assertAllClose(new_output_data, expected_output) if __name__ == "__main__": diff --git a/keras/layers/preprocessing/hashing_test.py b/keras/layers/preprocessing/hashing_test.py index 76f20719f6e..7bb20dc1eab 100644 --- a/keras/layers/preprocessing/hashing_test.py +++ b/keras/layers/preprocessing/hashing_test.py @@ -414,6 +414,30 @@ def test_saved_model(self): new_output_data = loaded_model(input_data) self.assertAllClose(new_output_data, original_output_data) + @test_utils.run_v2_only + def test_save_keras_v3(self): + input_data = np.array( + ["omar", "stringer", "marlo", "wire", "skywalker"] + ) + + inputs = keras.Input(shape=(None,), dtype=tf.string) + outputs = hashing.Hashing(num_bins=100)(inputs) + model = keras.Model(inputs=inputs, outputs=outputs) + + original_output_data = model(input_data) + + # Save the model to disk. + output_path = os.path.join(self.get_temp_dir(), "tf_keras_model.keras") + model.save(output_path, save_format="keras_v3") + loaded_model = keras.models.load_model(output_path) + + # Ensure that the loaded model is unique (so that the save/load is real) + self.assertIsNot(model, loaded_model) + + # Validate correctness of the new model. + new_output_data = loaded_model(input_data) + self.assertAllClose(new_output_data, original_output_data) + @parameterized.named_parameters( ( "list_input", diff --git a/keras/layers/preprocessing/index_lookup_test.py b/keras/layers/preprocessing/index_lookup_test.py index 91a8fc8b771..ca488eb4c54 100644 --- a/keras/layers/preprocessing/index_lookup_test.py +++ b/keras/layers/preprocessing/index_lookup_test.py @@ -2211,6 +2211,7 @@ def test_vocabulary_persistence_across_saving(self): ] ) expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]] + vocab_file = self._write_to_temp_file("temp", vocab_data) # Build and validate a golden model. input_data = keras.Input(shape=(None,), dtype=tf.string) @@ -2220,32 +2221,57 @@ def test_vocabulary_persistence_across_saving(self): mask_token="", oov_token="[OOV]", vocabulary_dtype=tf.string, + vocabulary=vocab_file, ) - layer.set_vocabulary(vocab_data) int_data = layer(input_data) model = keras.Model(inputs=input_data, outputs=int_data) output_dataset = model.predict(input_array) self.assertAllEqual(output_dataset, expected_output) - # Save the model to disk. - output_path = os.path.join(self.get_temp_dir(), "tf_keras_saved_model") - model.save(output_path, save_format="tf") + with self.subTest("keras_v3"): + # Save the model to disk. + output_path = os.path.join( + self.get_temp_dir(), "tf_keras_model.keras" + ) + model.save(output_path, save_format="keras_v3") - # Delete the session and graph to ensure that the loaded model is - # generated from scratch. - keras.backend.clear_session() + loaded_model = keras.models.load_model( + output_path, + custom_objects={"IndexLookup": index_lookup.IndexLookup}, + ) - loaded_model = keras.models.load_model( - output_path, - custom_objects={"IndexLookup": index_lookup.IndexLookup}, - ) + # Ensure that the loaded model is unique + # (so that the save/load is real) + self.assertIsNot(model, loaded_model) - # Ensure that the loaded model is unique (so that the save/load is real) - self.assertIsNot(model, loaded_model) + # Validate correctness of the new model. + new_output_dataset = loaded_model.predict(input_array) + self.assertAllEqual(new_output_dataset, expected_output) - # Validate correctness of the new model. - new_output_dataset = loaded_model.predict(input_array) - self.assertAllEqual(new_output_dataset, expected_output) + with self.subTest("savedmodel"): + # Save the model to disk. + output_path = os.path.join( + self.get_temp_dir(), "tf_keras_saved_model" + ) + model.save(output_path, save_format="tf") + + # Delete the session and graph to ensure that the loaded model is + # generated from scratch. + keras.backend.clear_session() + tf.io.gfile.remove(vocab_file) + + loaded_model = keras.models.load_model( + output_path, + custom_objects={"IndexLookup": index_lookup.IndexLookup}, + ) + + # Ensure that the loaded model is unique + # (so that the save/load is real) + self.assertIsNot(model, loaded_model) + + # Validate correctness of the new model. + new_output_dataset = loaded_model.predict(input_array) + self.assertAllEqual(new_output_dataset, expected_output) def test_vocabulary_persistence_file_across_cloning(self): vocab_data = ["earth", "wind", "and", "fire"] @@ -2401,56 +2427,108 @@ def test_persistence_file_vocab_keras_save_keras_load(self): output_dataset = model.predict(input_array) self.assertAllEqual(output_dataset, expected_output) - # Save the model to disk. - output_path = os.path.join(self.get_temp_dir(), "tf_keras_saved_model") - model.save(output_path, save_format="tf") + with self.subTest("keras_v3"): + # Save the model to disk. + output_path = os.path.join( + self.get_temp_dir(), "tf_keras_model.keras" + ) + model.save(output_path, save_format="keras_v3") - # Delete the session and graph to ensure that the loaded model is - # generated from scratch. - keras.backend.clear_session() - tf.io.gfile.remove(vocab_file) + loaded_model = keras.models.load_model( + output_path, + custom_objects={"IndexLookup": index_lookup.IndexLookup}, + ) - loaded_model = keras.models.load_model( - output_path, - custom_objects={"IndexLookup": index_lookup.IndexLookup}, - ) + # Ensure that the loaded model is unique + # (so that the save/load is real) + self.assertIsNot(model, loaded_model) + + # Validate correctness of the new model. + new_output_dataset = loaded_model.predict(input_array) + self.assertAllEqual(new_output_dataset, expected_output) + + # Try re-saving the layer. This simulates saving a layer + # contained at a hub Module. + input_data_2 = keras.Input(shape=(None,), dtype=tf.string) + output_2 = loaded_model(input_data_2) + model_2 = keras.Model(inputs=input_data_2, outputs=output_2) + new_output_dataset = model_2.predict(input_array) + self.assertAllEqual(new_output_dataset, expected_output) + + # Save the model to disk. + output_path = os.path.join( + self.get_temp_dir(), "tf_keras_model_2.keras" + ) + model_2.save(output_path, save_format="keras_v3") - # Ensure that the loaded model is unique (so that the save/load is real) - self.assertIsNot(model, loaded_model) + loaded_model = keras.models.load_model( + output_path, + custom_objects={"IndexLookup": index_lookup.IndexLookup}, + ) - # Validate correctness of the new model. - new_output_dataset = loaded_model.predict(input_array) - self.assertAllEqual(new_output_dataset, expected_output) + # Ensure that the loaded model is unique + # (so that the save/load is real) + self.assertIsNot(model, loaded_model) - # Try re-saving the layer. This simulates saving a layer contained at - # a hub Module. - input_data_2 = keras.Input(shape=(None,), dtype=tf.string) - output_2 = loaded_model(input_data_2) - model_2 = keras.Model(inputs=input_data_2, outputs=output_2) - new_output_dataset = model_2.predict(input_array) - self.assertAllEqual(new_output_dataset, expected_output) + # Validate correctness of the new model. + new_output_dataset = loaded_model.predict(input_array) + self.assertAllEqual(new_output_dataset, expected_output) - # Save the model to disk. - output_path = os.path.join( - self.get_temp_dir(), "tf_keras_saved_model_2" - ) - model_2.save(output_path, save_format="tf") + with self.subTest("saved_model"): + # Save the model to disk. + output_path = os.path.join( + self.get_temp_dir(), "tf_keras_saved_model" + ) + model.save(output_path, save_format="tf") - # Delete the session and graph to ensure that the loaded model is - # generated from scratch. - keras.backend.clear_session() + # Delete the session and graph to ensure that the loaded model is + # generated from scratch. + keras.backend.clear_session() + tf.io.gfile.remove(vocab_file) - loaded_model = keras.models.load_model( - output_path, - custom_objects={"IndexLookup": index_lookup.IndexLookup}, - ) + loaded_model = keras.models.load_model( + output_path, + custom_objects={"IndexLookup": index_lookup.IndexLookup}, + ) - # Ensure that the loaded model is unique (so that the save/load is real) - self.assertIsNot(model, loaded_model) + # Ensure that the loaded model is unique + # (so that the save/load is real) + self.assertIsNot(model, loaded_model) + + # Validate correctness of the new model. + new_output_dataset = loaded_model.predict(input_array) + self.assertAllEqual(new_output_dataset, expected_output) + + # Try re-saving the layer. This simulates saving a layer + # contained at a hub Module. + input_data_2 = keras.Input(shape=(None,), dtype=tf.string) + output_2 = loaded_model(input_data_2) + model_2 = keras.Model(inputs=input_data_2, outputs=output_2) + new_output_dataset = model_2.predict(input_array) + self.assertAllEqual(new_output_dataset, expected_output) + + # Save the model to disk. + output_path = os.path.join( + self.get_temp_dir(), "tf_keras_saved_model_2" + ) + model_2.save(output_path, save_format="tf") - # Validate correctness of the new model. - new_output_dataset = loaded_model.predict(input_array) - self.assertAllEqual(new_output_dataset, expected_output) + # Delete the session and graph to ensure that the loaded model is + # generated from scratch. + keras.backend.clear_session() + + loaded_model = keras.models.load_model( + output_path, + custom_objects={"IndexLookup": index_lookup.IndexLookup}, + ) + + # Ensure that the loaded model is unique + # (so that the save/load is real) + self.assertIsNot(model, loaded_model) + + # Validate correctness of the new model. + new_output_dataset = loaded_model.predict(input_array) + self.assertAllEqual(new_output_dataset, expected_output) def test_persistence_file_vocab_keras_save_keras_load_tf_save_tf_load(self): vocab_data = ["earth", "wind", "and", "fire"] diff --git a/keras/layers/preprocessing/integer_lookup_test.py b/keras/layers/preprocessing/integer_lookup_test.py index a99075db4d6..4a06475880c 100644 --- a/keras/layers/preprocessing/integer_lookup_test.py +++ b/keras/layers/preprocessing/integer_lookup_test.py @@ -630,27 +630,56 @@ def test_vocabulary_persistence_across_saving(self): output_dataset = model.predict(input_array) self.assertAllEqual(output_dataset, expected_output) - # Save the model to disk. - output_path = os.path.join(self.get_temp_dir(), "tf_keras_saved_model") - model.save(output_path, save_format="tf") - - # Delete the session and graph to ensure that the loaded model is - # generated from scratch. - # TODO(b/149526183): Can't clear session when TF2 is disabled. - if tf.__internal__.tf2.enabled(): - keras.backend.clear_session() - - loaded_model = keras.models.load_model( - output_path, - custom_objects={"IntegerLookup": integer_lookup.IntegerLookup}, - ) + with self.subTest("keras_v3"): + if not tf.__internal__.tf2.enabled(): + self.skipTest( + "TF2 must be enabled to use the new `.keras` saving." + ) + + # Save the model to disk. + output_path = os.path.join( + self.get_temp_dir(), "tf_keras_model.keras" + ) + model.save(output_path, save_format="keras_v3") + + loaded_model = keras.models.load_model( + output_path, + custom_objects={"IntegerLookup": integer_lookup.IntegerLookup}, + ) + + # Ensure that the loaded model is unique + # (so that the save/load is real) + self.assertIsNot(model, loaded_model) + + # Validate correctness of the new model. + new_output_dataset = loaded_model.predict(input_array) + self.assertAllEqual(new_output_dataset, expected_output) + + with self.subTest("savedmodel"): + # Save the model to disk. + output_path = os.path.join( + self.get_temp_dir(), "tf_keras_saved_model" + ) + model.save(output_path, save_format="tf") + + # Delete the session and graph to ensure that the loaded model is + # generated from scratch. + # TODO(b/149526183): Can't clear session when TF2 is disabled. + if tf.__internal__.tf2.enabled(): + keras.backend.clear_session() + + loaded_model = keras.models.load_model( + output_path, + custom_objects={"IntegerLookup": integer_lookup.IntegerLookup}, + ) - # Ensure that the loaded model is unique (so that the save/load is real) - self.assertIsNot(model, loaded_model) + # Ensure that the loaded model is unique + # (so that the save/load is real) + self.assertIsNot(model, loaded_model) - # Validate correctness of the new model. - new_output_dataset = loaded_model.predict(input_array) - self.assertAllEqual(new_output_dataset, expected_output) + # Validate correctness of the new model. + new_output_dataset = loaded_model.predict(input_array) + self.assertAllEqual(new_output_dataset, expected_output) if __name__ == "__main__": diff --git a/keras/layers/preprocessing/normalization_test.py b/keras/layers/preprocessing/normalization_test.py index c0ffdb26fa8..d948f34d38f 100644 --- a/keras/layers/preprocessing/normalization_test.py +++ b/keras/layers/preprocessing/normalization_test.py @@ -392,7 +392,7 @@ def test_multiple_adapts(self): {"adapted": True}, {"adapted": False}, ) - def test_saved_model_tf(self, adapted): + def test_saving_tf(self, adapted): input_data = [[0.0], [2.0], [0.0], [2.0]] expected_output = [[-1.0], [1.0], [-1.0], [1.0]] @@ -422,10 +422,10 @@ def test_saved_model_tf(self, adapted): self.assertAllClose(new_output_data, expected_output) @parameterized.product( - save_format=["tf", "h5"], + save_format=["tf", "h5", "keras_v3"], adapt=[True, False], ) - def test_saved_model_keras(self, save_format, adapt): + def test_saving_keras(self, save_format, adapt): input_data = [[0.0], [2.0], [0.0], [2.0]] expected_output = [[-1.0], [1.0], [-1.0], [1.0]] @@ -443,7 +443,13 @@ def test_saved_model_keras(self, save_format, adapt): self.assertAllClose(output_data, expected_output) # Save the model to disk. - output_path = os.path.join(self.get_temp_dir(), "tf_keras_saved_model") + output_path = os.path.join(self.get_temp_dir(), "tf_keras_model") + if save_format == "keras_v3": + if not tf.__internal__.tf2.enabled(): + self.skipTest( + "TF2 must be enabled to use the new `.keras` saving." + ) + output_path += ".keras" model.save(output_path, save_format=save_format) loaded_model = keras.models.load_model( output_path, custom_objects={"Normalization": cls} @@ -457,10 +463,10 @@ def test_saved_model_keras(self, save_format, adapt): self.assertAllClose(new_output_data, expected_output) @parameterized.product( - save_format=["tf", "h5"], + save_format=["tf", "h5", "keras_v3"], adapt=[True, False], ) - def test_saved_model_keras_invert(self, save_format, adapt): + def test_saving_keras_invert(self, save_format, adapt): expected_output = [[0.0], [2.0], [0.0], [2.0]] input_data = [[-1.0], [1.0], [-1.0], [1.0]] @@ -478,9 +484,13 @@ def test_saved_model_keras_invert(self, save_format, adapt): self.assertAllClose(output_data, expected_output) # Save the model to disk. - output_path = os.path.join( - self.get_temp_dir(), "tf_keras_saved_model_invert" - ) + output_path = os.path.join(self.get_temp_dir(), "tf_keras_model_invert") + if save_format == "keras_v3": + if not tf.__internal__.tf2.enabled(): + self.skipTest( + "TF2 must be enabled to use the new `.keras` saving." + ) + output_path += ".keras" model.save(output_path, save_format=save_format) loaded_model = keras.models.load_model( output_path, custom_objects={"Normalization": cls} diff --git a/keras/layers/regularization/dropout_test.py b/keras/layers/regularization/dropout_test.py index bf53b4a44ad..2239338b8af 100644 --- a/keras/layers/regularization/dropout_test.py +++ b/keras/layers/regularization/dropout_test.py @@ -67,7 +67,7 @@ def test_dropout_with_zero_rate(self): rng_state_var, dropout._random_generator._generator._state_var ) - def test_dropout_with_savemodel(self): + def test_dropout_with_saving(self): inputs = keras.Input(shape=(5, 10)) layer = keras.layers.Dropout(0.5, force_generator=True) outputs = layer(inputs) @@ -83,45 +83,68 @@ def test_dropout_with_savemodel(self): # Make sure the layer does dropout value when training self.assertNotAllClose(train, predict) - model.save( - os.path.join(self.get_temp_dir(), "savedmodel"), save_format="tf" - ) - loaded_model = keras.models.load_model( - os.path.join(self.get_temp_dir(), "savedmodel") - ) - predict2 = loaded_model(np.ones((20, 5, 10))) - - self.assertAllClose(predict, predict2) - # Make sure the model dropout different value after loading - train2 = loaded_model(np.ones((20, 5, 10)), training=True) - self.assertNotAllClose(train, train2) - self.assertIsNotNone(loaded_model.layers[1]._random_generator) - - # Also make sure the checkpoint doesn't contain any variable from the - # dropout layer, to keep the backward compatibility. - checkpoint = tf.train.Checkpoint(model) - save_path = checkpoint.save( - os.path.join(self.get_temp_dir(), "checkpoint") - ) - checkpoint_var_names = [ - name_value_tuple[0] - for name_value_tuple in tf.train.list_variables(save_path) - ] - for name in checkpoint_var_names: - self.assertNotIn("dropout", name) - - # Make sure the checkpoint can be loaded - clone_model = keras.models.clone_model(model) - checkpoint = tf.train.Checkpoint(clone_model) - status = checkpoint.restore( - os.path.join(self.get_temp_dir(), "checkpoint-1") - ) - self.assertTrue(status.assert_consumed()) - self.assertTrue(status.assert_existing_objects_matched()) - # Make sure the output is differnt from the original model, since - # the StateVar is not preserved. - train3 = clone_model(np.ones((20, 5, 10)), training=True) - self.assertNotAllClose(train3, train2) + with self.subTest("savedmodel"): + model.save( + os.path.join(self.get_temp_dir(), "savedmodel"), + save_format="tf", + ) + loaded_model = keras.models.load_model( + os.path.join(self.get_temp_dir(), "savedmodel") + ) + predict2 = loaded_model(np.ones((20, 5, 10))) + + self.assertAllClose(predict, predict2) + # Make sure the model dropout different value after loading + train2 = loaded_model(np.ones((20, 5, 10)), training=True) + self.assertNotAllClose(train, train2) + self.assertIsNotNone(loaded_model.layers[1]._random_generator) + + with self.subTest("keras_v3"): + if not tf.__internal__.tf2.enabled(): + self.skipTest( + "TF2 must be enabled to use the new `.keras` saving." + ) + model.save( + os.path.join(self.get_temp_dir(), "model.keras"), + save_format="keras_v3", + ) + loaded_model = keras.models.load_model( + os.path.join(self.get_temp_dir(), "model.keras") + ) + predict2 = loaded_model(np.ones((20, 5, 10))) + + self.assertAllClose(predict, predict2) + # Make sure the model dropout different value after loading + train2 = loaded_model(np.ones((20, 5, 10)), training=True) + self.assertNotAllClose(train, train2) + self.assertIsNotNone(loaded_model.layers[1]._random_generator) + + with self.subTest("checkpoint"): + # Also make sure the checkpoint doesn't contain any variable from + # the dropout layer, to keep the backward compatibility. + checkpoint = tf.train.Checkpoint(model) + save_path = checkpoint.save( + os.path.join(self.get_temp_dir(), "checkpoint") + ) + checkpoint_var_names = [ + name_value_tuple[0] + for name_value_tuple in tf.train.list_variables(save_path) + ] + for name in checkpoint_var_names: + self.assertNotIn("dropout", name) + + # Make sure the checkpoint can be loaded + clone_model = keras.models.clone_model(model) + checkpoint = tf.train.Checkpoint(clone_model) + status = checkpoint.restore( + os.path.join(self.get_temp_dir(), "checkpoint-1") + ) + self.assertTrue(status.assert_consumed()) + self.assertTrue(status.assert_existing_objects_matched()) + # Make sure the output is differnt from the original model, since + # the StateVar is not preserved. + train3 = clone_model(np.ones((20, 5, 10)), training=True) + self.assertNotAllClose(train3, train2) @test_utils.run_v2_only def test_state_variable_name(self):