fix quantization save and load error

Author: divyashreepathihalli

keras/src/layers/layer.py

@@ -1359,10 +1359,32 @@ def save_own_variables(self, store):
         Args:
             store: Dict where the state of the model will be saved.
         """
+        if not getattr(self, "_is_quantized", False):
+            all_vars = self._trainable_variables + self._non_trainable_variables
+            for i, v in enumerate(all_vars):
+                store[f"{i}"] = v
+            return
+
+        # Case: quantized layer
+        quantized_vars = self._get_quantized_variables()
+        for i, v in enumerate(quantized_vars):
+            store[f"quantized_{i}"] = v
+
+        # Save non-quantized variables
         all_vars = self._trainable_variables + self._non_trainable_variables
-        for i, v in enumerate(all_vars):
+        non_quantized_vars = [
+            v for v in all_vars if v not in quantized_vars and v.trainable
+        ]
+        for i, v in enumerate(non_quantized_vars):
             store[f"{i}"] = v
 
+    def _get_quantized_variables(self):
+        quantized_vars = []
+        for v in self._trainable_variables + self._non_trainable_variables:
+            if not backend.is_float_dtype(v.dtype):
+                quantized_vars.append(v)
+        return quantized_vars
+
     def load_own_variables(self, store):
         """Loads the state of the layer.
 
@@ -1372,6 +1394,10 @@ def load_own_variables(self, store):
         Args:
             store: Dict from which the state of the model will be loaded.
         """
+        if any(key.startswith("quantized_") for key in store.keys()):
+            self._load_quantized_variables(store)
+            return
+
         all_vars = self._trainable_variables + self._non_trainable_variables
         if len(store.keys()) != len(all_vars):
             if len(all_vars) == 0 and not self.built:
@@ -1407,6 +1433,19 @@ def load_own_variables(self, store):
         for i, v in enumerate(all_vars):
             v.assign(store[f"{i}"])
 
+    def _load_quantized_variables(self, store):
+        quantized_vars = self._get_quantized_variables()
+        for i, v in enumerate(quantized_vars):
+            v.assign(store[f"quantized_{i}"])
+
+        # Load non-quantized variables
+        all_vars = self._trainable_variables + self._non_trainable_variables
+        non_quantized_vars = [
+            v for v in all_vars if v not in quantized_vars and v.trainable
+        ]
+        for i, v in enumerate(non_quantized_vars):
+            v.assign(store[f"{i}"])
+
     def _track_variable(self, variable):
         if variable.trainable:
             self._tracker.add_to_store("trainable_variables", variable)

keras/src/layers/layer_test.py

@@ -1,3 +1,4 @@
+import os
 import pickle
 from unittest import mock
 
@@ -12,6 +13,7 @@
 from keras.src import metrics
 from keras.src import models
 from keras.src import ops
+from keras.src import saving
 from keras.src import testing
 from keras.src.backend.common import global_state
 from keras.src.backend.common.remat import RematScope
@@ -1758,3 +1760,16 @@ def call(self, x):
         # foo_mode omitted -> foo_mode defaults to False -> no change
         y2 = model(sample_input)
         self.assertAllClose(y2, sample_input)
+
+    def test_quantized_model_save_and_load(self):
+        inputs = layers.Input(shape=(None,))
+        x = layers.Embedding(input_dim=10, output_dim=10)(inputs)
+        x = layers.Dense(10)(x)
+        model = models.Model(inputs=inputs, outputs=x)
+        path = os.path.join(self.get_temp_dir(), "quantized_model.keras")
+        model.quantize(mode="int8")
+        model.save(path)
+
+        quantized_model = saving.load_model(path)
+
+        self.assertTrue(quantized_model.built)