Increase test coverage + Fix save_model_to_hdf5 + Improve is_remote_path + Fix is_remote_path

keras_core/callbacks/lambda_callback_test.py

@@ -13,6 +13,7 @@
 class LambdaCallbackTest(testing.TestCase):
     @pytest.mark.requires_trainable_backend
     def test_LambdaCallback(self):
+        """Test standard LambdaCallback functionalities with training."""
         BATCH_SIZE = 4
         model = Sequential(
             [layers.Input(shape=(2,), batch_size=BATCH_SIZE), layers.Dense(1)]
@@ -40,7 +41,129 @@ def test_LambdaCallback(self):
                 epochs=5,
                 verbose=0,
             )
-            self.assertTrue(any("on_train_begin" in log for log in logs.output))
-            self.assertTrue(any("on_epoch_begin" in log for log in logs.output))
-            self.assertTrue(any("on_epoch_end" in log for log in logs.output))
-            self.assertTrue(any("on_train_end" in log for log in logs.output))
+            self.assertTrue
+            (any("on_train_begin" in log for log in logs.output))
+            self.assertTrue
+            (any("on_epoch_begin" in log for log in logs.output))
+            self.assertTrue
+            (any("on_epoch_end" in log for log in logs.output))
+            self.assertTrue
+            (any("on_train_end" in log for log in logs.output))
+
+    @pytest.mark.requires_trainable_backend
+    def test_LambdaCallback_with_batches(self):
+        """Test LambdaCallback's behavior with batch-level callbacks."""
+        BATCH_SIZE = 4
+        model = Sequential(
+            [layers.Input(shape=(2,), batch_size=BATCH_SIZE), layers.Dense(1)]
+        )
+        model.compile(
+            optimizer=optimizers.SGD(), loss=losses.MeanSquaredError()
+        )
+        x = np.random.randn(16, 2)
+        y = np.random.randn(16, 1)
+        lambda_log_callback = callbacks.LambdaCallback(
+            on_train_batch_begin=lambda batch, logs: logging.warning(
+                "on_train_batch_begin"
+            ),
+            on_train_batch_end=lambda batch, logs: logging.warning(
+                "on_train_batch_end"
+            ),
+        )
+        with self.assertLogs(level="WARNING") as logs:
+            model.fit(
+                x,
+                y,
+                batch_size=BATCH_SIZE,
+                validation_split=0.2,
+                callbacks=[lambda_log_callback],
+                epochs=5,
+                verbose=0,
+            )
+            self.assertTrue(
+                any("on_train_batch_begin" in log for log in logs.output)
+            )
+            self.assertTrue(
+                any("on_train_batch_end" in log for log in logs.output)
+            )
+
+    @pytest.mark.requires_trainable_backend
+    def test_LambdaCallback_with_kwargs(self):
+        """Test LambdaCallback's behavior with custom defined callback."""
+        BATCH_SIZE = 4
+        model = Sequential(
+            [layers.Input(shape=(2,), batch_size=BATCH_SIZE), layers.Dense(1)]
+        )
+        model.compile(
+            optimizer=optimizers.SGD(), loss=losses.MeanSquaredError()
+        )
+        x = np.random.randn(16, 2)
+        y = np.random.randn(16, 1)
+        model.fit(
+            x, y, batch_size=BATCH_SIZE, epochs=1, verbose=0
+        )  # Train briefly for evaluation to work.
+
+        def custom_on_test_begin(logs):
+            logging.warning("custom_on_test_begin_executed")
+
+        lambda_log_callback = callbacks.LambdaCallback(
+            on_test_begin=custom_on_test_begin
+        )
+        with self.assertLogs(level="WARNING") as logs:
+            model.evaluate(
+                x,
+                y,
+                batch_size=BATCH_SIZE,
+                callbacks=[lambda_log_callback],
+                verbose=0,
+            )
+            self.assertTrue(
+                any(
+                    "custom_on_test_begin_executed" in log
+                    for log in logs.output
+                )
+            )
+
+    @pytest.mark.requires_trainable_backend
+    def test_LambdaCallback_no_args(self):
+        """Test initializing LambdaCallback without any arguments."""
+        lambda_callback = callbacks.LambdaCallback()
+        self.assertIsInstance(lambda_callback, callbacks.LambdaCallback)
+
+    @pytest.mark.requires_trainable_backend
+    def test_LambdaCallback_with_additional_kwargs(self):
+        """Test initializing LambdaCallback with non-predefined kwargs."""
+
+        def custom_callback(logs):
+            pass
+
+        lambda_callback = callbacks.LambdaCallback(
+            custom_method=custom_callback
+        )
+        self.assertTrue(hasattr(lambda_callback, "custom_method"))
+
+    @pytest.mark.requires_trainable_backend
+    def test_LambdaCallback_during_prediction(self):
+        """Test LambdaCallback's functionality during model prediction."""
+        BATCH_SIZE = 4
+        model = Sequential(
+            [layers.Input(shape=(2,), batch_size=BATCH_SIZE), layers.Dense(1)]
+        )
+        model.compile(
+            optimizer=optimizers.SGD(), loss=losses.MeanSquaredError()
+        )
+        x = np.random.randn(16, 2)
+
+        def custom_on_predict_begin(logs):
+            logging.warning("on_predict_begin_executed")
+
+        lambda_callback = callbacks.LambdaCallback(
+            on_predict_begin=custom_on_predict_begin
+        )
+        with self.assertLogs(level="WARNING") as logs:
+            model.predict(
+                x, batch_size=BATCH_SIZE, callbacks=[lambda_callback], verbose=0
+            )
+            self.assertTrue(
+                any("on_predict_begin_executed" in log for log in logs.output)
+            )

keras_core/legacy/saving/legacy_h5_format.py

@@ -37,10 +37,9 @@ def save_model_to_hdf5(model, filepath, overwrite=True, include_optimizer=True):
             if not proceed:
                 return
 
-        # Try creating dir if not exist
         dirpath = os.path.dirname(filepath)
-        if not os.path.exists(dirpath):
-            os.path.makedirs(dirpath)
+        if dirpath and not os.path.exists(dirpath):
+            os.makedirs(dirpath, exist_ok=True)
 
         f = h5py.File(filepath, mode="w")
         opened_new_file = True

keras_core/legacy/saving/legacy_h5_format_test.py

@@ -481,3 +481,19 @@ def call(self, x):
 
         # Compare output
         self.assertAllClose(ref_output, output, atol=1e-5)
+
+
+@pytest.mark.requires_trainable_backend
+class DirectoryCreationTest(testing.TestCase):
+    def test_directory_creation_on_save(self):
+        """Test if directory is created on model save."""
+        model = get_sequential_model(keras_core)
+        nested_dirpath = os.path.join(
+            self.get_temp_dir(), "dir1", "dir2", "dir3"
+        )
+        filepath = os.path.join(nested_dirpath, "model.h5")
+        self.assertFalse(os.path.exists(nested_dirpath))
+        legacy_h5_format.save_model_to_hdf5(model, filepath)
+        self.assertTrue(os.path.exists(nested_dirpath))
+        loaded_model = legacy_h5_format.load_model_from_hdf5(filepath)
+        self.assertEqual(model.to_json(), loaded_model.to_json())

keras_core/saving/saving_api_test.py

@@ -0,0 +1,190 @@
+import os
+import unittest.mock as mock
+
+import numpy as np
+from absl import logging
+
+from keras_core import layers
+from keras_core.models import Sequential
+from keras_core.saving import saving_api
+from keras_core.testing import test_case
+
+
+class SaveModelTests(test_case.TestCase):
+    def setUp(self):
+        self.model = Sequential(
+            [
+                layers.Dense(5, input_shape=(3,)),
+                layers.Softmax(),
+            ],
+        )
+        self.filepath = "test_model.keras"
+        if os.path.exists(self.filepath):
+            os.remove(self.filepath)
+        saving_api.save_model(self.model, self.filepath)
+
+    def test_basic_saving(self):
+        """Test basic model saving and loading."""
+        loaded_model = saving_api.load_model(self.filepath)
+        x = np.random.uniform(size=(10, 3))
+        self.assertTrue(
+            np.allclose(self.model.predict(x), loaded_model.predict(x))
+        )
+
+    def test_invalid_save_format(self):
+        """Test deprecated save_format argument."""
+        with self.assertRaisesRegex(
+            ValueError, "The `save_format` argument is deprecated"
+        ):
+            saving_api.save_model(self.model, "model.txt", save_format=True)
+
+    def test_unsupported_arguments(self):
+        """Test unsupported argument during model save."""
+        with self.assertRaisesRegex(
+            ValueError, r"The following argument\(s\) are not supported"
+        ):
+            saving_api.save_model(self.model, self.filepath, random_arg=True)
+
+    def test_save_h5_format(self):
+        """Test saving model in h5 format."""
+        filepath_h5 = "test_model.h5"
+        saving_api.save_model(self.model, filepath_h5)
+        self.assertTrue(os.path.exists(filepath_h5))
+        os.remove(filepath_h5)
+
+    def test_save_unsupported_extension(self):
+        """Test saving model with unsupported extension."""
+        with self.assertRaisesRegex(
+            ValueError, "Invalid filepath extension for saving"
+        ):
+            saving_api.save_model(self.model, "model.png")
+
+    def tearDown(self):
+        if os.path.exists(self.filepath):
+            os.remove(self.filepath)
+
+
+class LoadModelTests(test_case.TestCase):
+    def setUp(self):
+        self.model = Sequential(
+            [
+                layers.Dense(5, input_shape=(3,)),
+                layers.Softmax(),
+            ],
+        )
+        self.filepath = "test_model.keras"
+        saving_api.save_model(self.model, self.filepath)
+
+    def test_basic_load(self):
+        """Test basic model loading."""
+        loaded_model = saving_api.load_model(self.filepath)
+        x = np.random.uniform(size=(10, 3))
+        self.assertTrue(
+            np.allclose(self.model.predict(x), loaded_model.predict(x))
+        )
+
+    def test_load_unsupported_format(self):
+        """Test loading model with unsupported format."""
+        with self.assertRaisesRegex(ValueError, "File format not supported"):
+            saving_api.load_model("model.pkl")
+
+    def test_load_keras_not_zip(self):
+        """Test loading keras file that's not a zip."""
+        with self.assertRaisesRegex(ValueError, "File not found"):
+            saving_api.load_model("not_a_zip.keras")
+
+    def test_load_h5_format(self):
+        """Test loading model in h5 format."""
+        filepath_h5 = "test_model.h5"
+        saving_api.save_model(self.model, filepath_h5)
+        loaded_model = saving_api.load_model(filepath_h5)
+        x = np.random.uniform(size=(10, 3))
+        self.assertTrue(
+            np.allclose(self.model.predict(x), loaded_model.predict(x))
+        )
+        os.remove(filepath_h5)
+
+    def test_load_model_with_custom_objects(self):
+        """Test loading model with custom objects."""
+
+        class CustomLayer(layers.Layer):
+            def call(self, inputs):
+                return inputs
+
+        model = Sequential([CustomLayer(input_shape=(3,))])
+        model.save("custom_model.keras")
+        loaded_model = saving_api.load_model(
+            "custom_model.keras", custom_objects={"CustomLayer": CustomLayer}
+        )
+        self.assertIsInstance(loaded_model.layers[0], CustomLayer)
+        os.remove("custom_model.keras")
+
+    def tearDown(self):
+        if os.path.exists(self.filepath):
+            os.remove(self.filepath)
+
+
+class LoadWeightsTests(test_case.TestCase):
+    def setUp(self):
+        self.model = Sequential(
+            [
+                layers.Dense(5, input_shape=(3,)),
+                layers.Softmax(),
+            ],
+        )
+
+    def test_load_keras_weights(self):
+        """Test loading keras weights."""
+        filepath = "test_weights.weights.h5"
+        self.model.save_weights(filepath)
+        original_weights = self.model.get_weights()
+        self.model.load_weights(filepath)
+        loaded_weights = self.model.get_weights()
+        for orig, loaded in zip(original_weights, loaded_weights):
+            self.assertTrue(np.array_equal(orig, loaded))
+
+    def test_load_h5_weights_by_name(self):
+        """Test loading h5 weights by name."""
+        filepath = "test_weights.weights.h5"
+        self.model.save_weights(filepath)
+        with self.assertRaisesRegex(ValueError, "Invalid keyword arguments"):
+            self.model.load_weights(filepath, by_name=True)
+
+    def test_load_weights_invalid_extension(self):
+        """Test loading weights with unsupported extension."""
+        with self.assertRaisesRegex(ValueError, "File format not supported"):
+            self.model.load_weights("invalid_extension.pkl")
+
+    def tearDown(self):
+        filepath = "test_weights.weights.h5"
+        if os.path.exists(filepath):
+            os.remove(filepath)
+
+
+class SaveModelTestsWarning(test_case.TestCase):
+    def setUp(self):
+        self.model = Sequential(
+            [
+                layers.Dense(5, input_shape=(3,)),
+                layers.Softmax(),
+            ],
+        )
+        self.filepath = "test_model.keras"
+        if os.path.exists(self.filepath):
+            os.remove(self.filepath)
+        saving_api.save_model(self.model, self.filepath)
+
+    def test_h5_deprecation_warning(self):
+        """Test deprecation warning for h5 format."""
+        with mock.patch.object(logging, "warning") as mock_warn:
+            saving_api.save_model(self.model, "test_model.h5")
+            mock_warn.assert_called_once_with(
+                "You are saving your model as an HDF5 file via `model.save()`. "
+                "This file format is considered legacy. "
+                "We recommend using instead the native Keras format, "
+                "e.g. `model.save('my_model.keras')`."
+            )
+
+    def tearDown(self):
+        if os.path.exists(self.filepath):
+            os.remove(self.filepath)