Repository: keras-team/keras

Issue:
.keras model with base model trainable has poor performance in TensorFlow 2.17 (Keras 3)
I originally posted this bug in the [TensorFlow github issues section](https://github.com/tensorflow/tensorflow/issues/74170) since I believed it to be due to a higher TF version, but was asked to post here since it may instead be due to Keras 3. I am copying my post below:


I am training an EfficientNet model with a custom head using TensorFlow and Keras, saving the model to a .keras format. If the base model trainable flag is set to False, such that I only train the head, then when I later load the .keras model and evaluate it on a dataset, I get the expected good performance. When I set the trainable flag to True and train a model (which converges well), then when I later load the model and evaluate it on the same dataset the performance has degraded significantly. (I am evaluating the model on the same dataset using the same code both at the end of training, and later on in a separate notebook. It is in this separate notebook where the performance is bad, where again the same dataset is being used and the same code is being used in both evaluation places.)

Saving to a .h5 model does not have this issue, and the performance of the saved model is good. I have spent the day trying different trainable and training flag values in various places to no improvement, thinking originally that it was something to do with the BatchNorm layers in the model. Recompiling the model has not helped.

When I switch back to an older TensorFlow version (2.15.0.post1) with Keras 2 I do not see this issue. Both the trained .keras and .h5 models perform well when later loaded and evaluated on my dataset of interest.

This seems like a bug to me, though I also acknowledge that perhaps I have missed something in the TF/Keras updates. I have searched the TensorFlow API docs for the various methods to no success. If it is the latter I would be very grateful for any advice, thank you.


Discussion:
It seems there is a bug in the `EfficientNet` family.

Here is the reproducible script:

```python
import numpy as np

import keras

# Build custom model from EfficientNetB0.
img_size = 224
num_classes = 10
inputs = keras.Input((img_size, img_size, 3))
base_model = keras.applications.EfficientNetB0(
    weights="imagenet", include_top=False
)
x = base_model(inputs)
x = keras.layers.Dense(512, activation="relu")(x)
x = keras.layers.BatchNormalization()(x)
x = keras.layers.Dropout(rate=0.5)(x)

# Another fully connected dense layer
x = keras.layers.Dense(256, activation="relu")(x)
x = keras.layers.BatchNormalization()(x)
x = keras.layers.Dropout(rate=0.5)(x)

outputs = keras.layers.Dense(num_classes, activation="softmax")(x)
model = keras.Model(inputs, outputs)
weights = model.get_weights()

# Verify weights.
model.save("model.keras")
revived_model = keras.models.load_model("model.keras")
revived_weights = revived_model.get_weights()

np.testing.assert_equal(len(weights), len(revived_weights))
for w, rw in zip(weights, revived_weights):
    np.testing.assert_allclose(w, rw)
print("Weights are all the same.")

# Predict.
images = np.random.uniform(size=(2, 224, 224, 3)).astype("float32")
outputs = model.predict(images)
revived_outputs = revived_model.predict(images)
np.testing.assert_allclose(outputs, revived_outputs)  # This line failed.
print("Predictions are all the same.")

```

However, the script works fine with:
- `MobileNet*`
- `EfficientNetV2`

The root cause might be in `keras.applications.EfficientNet` since the saving / reloading work fine with other models.

Thanks for the report -- tried investigating this for a bit, but I can't figure it out for now. Interestingly the issue only occurs after loading ImageNet weights.



Patch:
diff --git a/keras/src/legacy/saving/legacy_h5_format.py b/keras/src/legacy/saving/legacy_h5_format.py
index 05f7ceca22be..6b6fa76af51e 100644
--- a/keras/src/legacy/saving/legacy_h5_format.py
+++ b/keras/src/legacy/saving/legacy_h5_format.py
@@ -372,8 +372,12 @@ def load_weights_from_hdf5_group(f, model):
                 f"Layer expects {len(symbolic_weights)} weight(s). Received "
                 f"{len(weight_values)} saved weight(s)"
             )
-        for ref_v, val in zip(symbolic_weights, weight_values):
-            ref_v.assign(val)
+        _set_weights(
+            layer,
+            symbolic_weights,
+            weight_values,
+            name=f"layer #{k} (named {layer.name})",
+        )
 
     if "top_level_model_weights" in f:
         symbolic_weights = list(
@@ -392,8 +396,57 @@ def load_weights_from_hdf5_group(f, model):
                 f"Model expects {len(symbolic_weights)} top-level weight(s). "
                 f"Received {len(weight_values)} saved top-level weight(s)"
             )
-        for ref_v, val in zip(symbolic_weights, weight_values):
-            ref_v.assign(val)
+        _set_weights(
+            model,
+            symbolic_weights,
+            weight_values,
+            name="top-level model",
+        )
+
+
+def _set_weights(
+    instance, symbolic_weights, weight_values, name, skip_mismatch=False
+):
+    """Safely set weights into a model or a layer,
+       calling load_own_variables if implemented.
+    Args:
+        instance: Model or layer instance,
+        symbolic_weights: symbolic tensors representing
+                        the weights of the variables to load,
+        weight_values: values of the weights to load,
+        skip_mismatch: Boolean, whether to skip loading of weights
+            where there is a mismatch in the shape of the weights,
+        name: name used to identify the group.
+    Raises:
+        ValueError: in case of mismatch between provided
+            model/layer and weights.
+    """
+    for i, weight_value in enumerate(weight_values):
+        expected_shape = symbolic_weights[i].shape
+        received_shape = weight_value.shape
+        if expected_shape != received_shape:
+            if skip_mismatch:
+                warnings.warn(
+                    f"Skipping loading weights for {name}"
+                    f"due to mismatch in shape for "
+                    f"weight {symbolic_weights[i].path}. "
+                    f"Weight expects shape {expected_shape}. "
+                    "Received saved weight "
+                    f"with shape {received_shape}",
+                    stacklevel=2,
+                )
+                continue
+            raise ValueError(
+                f"Shape mismatch in {name}"
+                f"for weight {symbolic_weights[i].path}. "
+                f"Weight expects shape {expected_shape}. "
+                "Received saved weight "
+                f"with shape {received_shape}"
+            )
+        symbolic_weights[i].assign(weight_value)
+
+    if hasattr(instance, "finalize_state") and symbolic_weights:
+        instance.finalize_state()
 
 
 def load_weights_from_hdf5_group_by_name(f, model, skip_mismatch=False):
@@ -456,30 +509,13 @@ def load_weights_from_hdf5_group_by_name(f, model, skip_mismatch=False):
                     f"Received {len(weight_values)} saved weight(s)"
                 )
             # Set values.
-            for i in range(len(weight_values)):
-                expected_shape = symbolic_weights[i].shape
-                received_shape = weight_values[i].shape
-                if expected_shape != received_shape:
-                    if skip_mismatch:
-                        warnings.warn(
-                            f"Skipping loading weights for layer #{k} (named "
-                            f"{layer.name}) due to mismatch in shape for "
-                            f"weight {symbolic_weights[i].path}. "
-                            f"Weight expects shape {expected_shape}. "
-                            "Received saved weight "
-                            f"with shape {received_shape}",
-                            stacklevel=2,
-                        )
-                        continue
-                    raise ValueError(
-                        f"Shape mismatch in layer #{k} (named {layer.name}) "
-                        f"for weight {symbolic_weights[i].path}. "
-                        f"Weight expects shape {expected_shape}. "
-                        "Received saved weight "
-                        f"with shape {received_shape}"
-                    )
-                else:
-                    symbolic_weights[i].assign(weight_values[i])
+            _set_weights(
+                layer,
+                symbolic_weights,
+                weight_values,
+                skip_mismatch=skip_mismatch,
+                name=f"layer #{k} (named {layer.name})",
+            )
 
     if "top_level_model_weights" in f:
         symbolic_weights = model.trainable_weights + model.non_trainable_weights
@@ -505,30 +541,13 @@ def load_weights_from_hdf5_group_by_name(f, model, skip_mismatch=False):
                     f"Received {len(weight_values)} saved top-level weight(s)"
                 )
         else:
-            for i in range(len(weight_values)):
-                expected_shape = symbolic_weights[i].shape
-                received_shape = weight_values[i].shape
-                if expected_shape != received_shape:
-                    if skip_mismatch:
-                        warnings.warn(
-                            "Skipping loading top-level weight for model due "
-                            "to mismatch in shape for "
-                            f"weight {symbolic_weights[i].path}. "
-                            f"Weight expects shape {expected_shape}. "
-                            "Received saved weight "
-                            f"with shape {received_shape}",
-                            stacklevel=2,
-                        )
-                    else:
-                        raise ValueError(
-                            "Shape mismatch in model for top-level weight "
-                            f"{symbolic_weights[i].path}. "
-                            f"Weight expects shape {expected_shape}. "
-                            "Received saved weight "
-                            f"with shape {received_shape}"
-                        )
-                else:
-                    symbolic_weights[i].assign(weight_values[i])
+            _set_weights(
+                model,
+                symbolic_weights,
+                weight_values,
+                skip_mismatch=skip_mismatch,
+                name="top-level model",
+            )
 
 
 def load_subset_weights_from_hdf5_group(f):


Test patch:
diff --git a/keras/src/saving/saving_lib_test.py b/keras/src/saving/saving_lib_test.py
index e308ba8919f5..15acea3287c3 100644
--- a/keras/src/saving/saving_lib_test.py
+++ b/keras/src/saving/saving_lib_test.py
@@ -959,6 +959,24 @@ def call(self, x):
             model_a.dense.kernel.numpy(), model_b.dense.kernel.numpy()
         )
 
+    def test_normalization_legacy_h5_format(self):
+        temp_filepath = os.path.join(self.get_temp_dir(), "custom_model.h5")
+
+        inputs = keras.Input((32,))
+        normalization = keras.layers.Normalization()
+        outputs = normalization(inputs)
+
+        model = keras.Model(inputs, outputs)
+
+        x = np.random.random((1, 32))
+        normalization.adapt(x)
+        ref_out = model(x)
+
+        model.save(temp_filepath)
+        new_model = keras.saving.load_model(temp_filepath)
+        out = new_model(x)
+        self.assertAllClose(ref_out, out, atol=1e-6)
+
     def test_legacy_h5_format(self):
         temp_filepath = os.path.join(self.get_temp_dir(), "custom_model.h5")