Repository: keras-team/keras
Issue category: Performance Issue

model.keras format much slower to load
Anyone experiencing unreasonably slow load times when loading a keras-format saved model? I have noticed this repeated when working in ipython, where simply instantiating a model via `Model.from_config` then calling `model.load_weights` is much (several factors) faster than loading a `model.keras` file.

My understanding is the keras format is simply a zip file with the config.json file and weights h5 (iirc) but weirdly enough, there's something not right going on while loading.


Functions to edit:
keras/src/saving/saving_lib.py:_save_model_to_fileobj
keras/src/saving/saving_lib.py:_load_model_from_fileobj
keras/src/saving/saving_lib.py:load_weights_only

Patch:
diff --git a/keras/src/saving/saving_lib.py b/keras/src/saving/saving_lib.py
index c16d2ffafb4..7e858025963 100644
--- a/keras/src/saving/saving_lib.py
+++ b/keras/src/saving/saving_lib.py
@@ -3,6 +3,7 @@
 import datetime
 import io
 import json
+import pathlib
 import tempfile
 import warnings
 import zipfile
@@ -32,6 +33,8 @@
 _CONFIG_FILENAME = "config.json"
 _METADATA_FILENAME = "metadata.json"
 _VARS_FNAME = "model.weights"  # Will become e.g. "model.weights.h5"
+_VARS_FNAME_H5 = _VARS_FNAME + ".h5"
+_VARS_FNAME_NPZ = _VARS_FNAME + ".npz"
 _ASSETS_DIRNAME = "assets"
 
 
@@ -109,30 +112,50 @@ def _save_model_to_fileobj(model, fileobj, weights_format):
         with zf.open(_CONFIG_FILENAME, "w") as f:
             f.write(config_json.encode())
 
-        if weights_format == "h5":
-            weights_store = H5IOStore(_VARS_FNAME + ".h5", archive=zf, mode="w")
-        elif weights_format == "npz":
-            weights_store = NpzIOStore(
-                _VARS_FNAME + ".npz", archive=zf, mode="w"
-            )
-        else:
-            raise ValueError(
-                "Unknown `weights_format` argument. "
-                "Expected 'h5' or 'npz'. "
-                f"Received: weights_format={weights_format}"
-            )
+        weights_file_path = None
+        try:
+            if weights_format == "h5":
+                if isinstance(fileobj, io.BufferedWriter):
+                    # First, open the .h5 file, then write it to `zf` at the end
+                    # of the function call.
+                    working_dir = pathlib.Path(fileobj.name).parent
+                    weights_file_path = working_dir / _VARS_FNAME_H5
+                    weights_store = H5IOStore(weights_file_path, mode="w")
+                else:
+                    # Fall back when `fileobj` is an `io.BytesIO`. Typically,
+                    # this usage is for pickling.
+                    weights_store = H5IOStore(
+                        _VARS_FNAME_H5, archive=zf, mode="w"
+                    )
+            elif weights_format == "npz":
+                weights_store = NpzIOStore(
+                    _VARS_FNAME_NPZ, archive=zf, mode="w"
+                )
+            else:
+                raise ValueError(
+                    "Unknown `weights_format` argument. "
+                    "Expected 'h5' or 'npz'. "
+                    f"Received: weights_format={weights_format}"
+                )
 
-        asset_store = DiskIOStore(_ASSETS_DIRNAME, archive=zf, mode="w")
+            asset_store = DiskIOStore(_ASSETS_DIRNAME, archive=zf, mode="w")
 
-        _save_state(
-            model,
-            weights_store=weights_store,
-            assets_store=asset_store,
-            inner_path="",
-            visited_saveables=set(),
-        )
-        weights_store.close()
-        asset_store.close()
+            _save_state(
+                model,
+                weights_store=weights_store,
+                assets_store=asset_store,
+                inner_path="",
+                visited_saveables=set(),
+            )
+        except:
+            # Skip the final `zf.write` if any exception is raised
+            weights_file_path = None
+            raise
+        finally:
+            weights_store.close()
+            asset_store.close()
+            if weights_file_path:
+                zf.write(weights_file_path, weights_file_path.name)
 
 
 def load_model(filepath, custom_objects=None, compile=True, safe_mode=True):
@@ -172,36 +195,49 @@ def _load_model_from_fileobj(fileobj, custom_objects, compile, safe_mode):
             )
 
         all_filenames = zf.namelist()
-        if _VARS_FNAME + ".h5" in all_filenames:
-            weights_store = H5IOStore(_VARS_FNAME + ".h5", archive=zf, mode="r")
-        elif _VARS_FNAME + ".npz" in all_filenames:
-            weights_store = NpzIOStore(
-                _VARS_FNAME + ".npz", archive=zf, mode="r"
-            )
-        else:
-            raise ValueError(
-                f"Expected a {_VARS_FNAME}.h5 or {_VARS_FNAME}.npz file."
-            )
+        weights_file_path = None
+        try:
+            if _VARS_FNAME_H5 in all_filenames:
+                if isinstance(fileobj, io.BufferedReader):
+                    # First, extract the model.weights.h5 file, then load it
+                    # using h5py.
+                    working_dir = pathlib.Path(fileobj.name).parent
+                    zf.extract(_VARS_FNAME_H5, working_dir)
+                    weights_file_path = working_dir / _VARS_FNAME_H5
+                    weights_store = H5IOStore(weights_file_path, mode="r")
+                else:
+                    # Fall back when `fileobj` is an `io.BytesIO`. Typically,
+                    # this usage is for pickling.
+                    weights_store = H5IOStore(_VARS_FNAME_H5, zf, mode="r")
+            elif _VARS_FNAME_NPZ in all_filenames:
+                weights_store = NpzIOStore(_VARS_FNAME_NPZ, zf, mode="r")
+            else:
+                raise ValueError(
+                    f"Expected a {_VARS_FNAME_H5} or {_VARS_FNAME_NPZ} file."
+                )
 
-        if len(all_filenames) > 3:
-            asset_store = DiskIOStore(_ASSETS_DIRNAME, archive=zf, mode="r")
-        else:
-            asset_store = None
-
-        failed_saveables = set()
-        error_msgs = {}
-        _load_state(
-            model,
-            weights_store=weights_store,
-            assets_store=asset_store,
-            inner_path="",
-            visited_saveables=set(),
-            failed_saveables=failed_saveables,
-            error_msgs=error_msgs,
-        )
-        weights_store.close()
-        if asset_store:
-            asset_store.close()
+            if len(all_filenames) > 3:
+                asset_store = DiskIOStore(_ASSETS_DIRNAME, archive=zf, mode="r")
+            else:
+                asset_store = None
+
+            failed_saveables = set()
+            error_msgs = {}
+            _load_state(
+                model,
+                weights_store=weights_store,
+                assets_store=asset_store,
+                inner_path="",
+                visited_saveables=set(),
+                failed_saveables=failed_saveables,
+                error_msgs=error_msgs,
+            )
+        finally:
+            weights_store.close()
+            if asset_store:
+                asset_store.close()
+            if weights_file_path:
+                weights_file_path.unlink()
 
         if failed_saveables:
             _raise_loading_failure(error_msgs)
@@ -250,9 +286,7 @@ def load_weights_only(
         weights_store = H5IOStore(filepath, mode="r")
     elif filepath.endswith(".keras"):
         archive = zipfile.ZipFile(filepath, "r")
-        weights_store = H5IOStore(
-            _VARS_FNAME + ".h5", archive=archive, mode="r"
-        )
+        weights_store = H5IOStore(_VARS_FNAME_H5, archive=archive, mode="r")
 
     failed_saveables = set()
     if objects_to_skip is not None:

Test patch: diff --git a/keras/src/saving/saving_lib_test.py b/keras/src/saving/saving_lib_test.py
index 80f53608f6c..91e40426285 100644
--- a/keras/src/saving/saving_lib_test.py
+++ b/keras/src/saving/saving_lib_test.py
@@ -612,6 +612,53 @@ def test_save_to_fileobj(self) -> None:
 
         self.assertAllClose(pred1, pred2, atol=1e-5)
 
+    def test_save_model_exception_raised(self):
+        # Assume we have an error in `save_own_variables`.
+        class RaiseErrorLayer(keras.layers.Layer):
+            def __init__(self, **kwargs):
+                super().__init__(**kwargs)
+
+            def call(self, inputs):
+                return inputs
+
+            def save_own_variables(self, store):
+                raise ValueError
+
+        model = keras.Sequential([keras.Input([1]), RaiseErrorLayer()])
+        filepath = f"{self.get_temp_dir()}/model.keras"
+        with self.assertRaises(ValueError):
+            saving_lib.save_model(model, filepath)
+
+        # Ensure we don't have a bad "model.weights.h5" inside the zip file.
+        self.assertTrue(Path(filepath).exists())
+        with zipfile.ZipFile(filepath) as zf:
+            all_filenames = zf.namelist()
+            self.assertNotIn("model.weights.h5", all_filenames)
+
+    def test_load_model_exception_raised(self):
+        # Assume we have an error in `load_own_variables`.
+        class RaiseErrorLayer(keras.layers.Layer):
+            def __init__(self, **kwargs):
+                super().__init__(**kwargs)
+
+            def call(self, inputs):
+                return inputs
+
+            def load_own_variables(self, store):
+                raise ValueError
+
+        model = keras.Sequential([keras.Input([1]), RaiseErrorLayer()])
+        filepath = f"{self.get_temp_dir()}/model.keras"
+        saving_lib.save_model(model, filepath)
+        with self.assertRaises(ValueError):
+            saving_lib.load_model(
+                filepath, custom_objects={"RaiseErrorLayer": RaiseErrorLayer}
+            )
+
+        # Ensure we don't leave a bad "model.weights.h5" on the filesystem.
+        self.assertTrue(Path(filepath).exists())
+        self.assertFalse(Path(filepath).with_name("model.weights.h5").exists())
+
 
 @pytest.mark.requires_trainable_backend
 class SavingAPITest(testing.TestCase):