Repository: keras-team/keras
Issue: [BUG] Model state is not correctly after saving and loading


**System information**.
- Have I written custom code (as opposed to using a stock example script provided in Keras): Yes
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux
- TensorFlow installed from (source or binary): binary
- TensorFlow version (use command below): 2.8.0
- Python version: 3.7.10
- Bazel version (if compiling from source):
- GPU model and memory: cpu
- Exact command to reproduce:

**Describe the problem**.

Keras model does not converge after saving and loading.

**Describe the current behavior**.
After calling `model.save(...)` and `model = tf.keras.models.load_model(...)`, the model failed to converge.

**Describe the expected behavior**.
Adding `model.save(...)` and `model = tf.keras.models.load_model(...)` should not effect the training process.

- Do you want to contribute a PR? (yes/no):  No

**Standalone code to reproduce the issue**.

```python
import tensorflow as tf
mnist = tf.keras.datasets.mnist

(x_train, y_train),(x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
########### added lines ###########
model.save("/tmp/mnist_model")
model = tf.keras.models.load_model('/tmp/mnist_model')
################################

model.fit(x_train, y_train, epochs=5)
model.evaluate(x_test, y_test)
```
**Source code / logs**.

After adding saving and loading, the model does not converge:
```
2022-02-08 09:03:19.403489: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2022-02-08 09:03:19.698194: W tensorflow/python/util/util.cc:368] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
Epoch 1/5
1875/1875 [==============================] - 6s 3ms/step - loss: 0.2983 - accuracy: 0.1004
Epoch 2/5
1875/1875 [==============================] - 6s 3ms/step - loss: 0.1451 - accuracy: 0.0992
Epoch 3/5
1875/1875 [==============================] - 6s 3ms/step - loss: 0.1069 - accuracy: 0.0990
Epoch 4/5
1875/1875 [==============================] - 6s 3ms/step - loss: 0.0881 - accuracy: 0.0991
Epoch 5/5
1875/1875 [==============================] - 6s 3ms/step - loss: 0.0754 - accuracy: 0.0989
313/313 [==============================] - 1s 2ms/step - loss: 0.0775 - accuracy: 0.0992
```

Remove the saving and loading code, the model converges as expected:

```
2022-02-08 09:05:18.683461: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
Epoch 1/5
1875/1875 [==============================] - 6s 3ms/step - loss: 0.3010 - accuracy: 0.9133   
Epoch 2/5
1875/1875 [==============================] - 6s 3ms/step - loss: 0.1452 - accuracy: 0.9574
Epoch 3/5
1875/1875 [==============================] - 6s 3ms/step - loss: 0.1099 - accuracy: 0.9666
Epoch 4/5
1875/1875 [==============================] - 6s 3ms/step - loss: 0.0878 - accuracy: 0.9729
Epoch 5/5
1875/1875 [==============================] - 6s 3ms/step - loss: 0.0746 - accuracy: 0.9766
313/313 [==============================] - 1s 2ms/step - loss: 0.0744 - accuracy: 0.9770
```


Changed files:
keras/saving/saving_utils.py

diff --git a/keras/saving/saving_utils.py b/keras/saving/saving_utils.py
index 6a6fa3b784..44b99e4533 100644
--- a/keras/saving/saving_utils.py
+++ b/keras/saving/saving_utils.py
@@ -317,12 +317,29 @@ def _enforce_names_consistency(specs):
 
 
 def try_build_compiled_arguments(model):
+  """Trying to build compiled metrics."""
   if (not version_utils.is_v1_layer_or_model(model) and
       model.outputs is not None):
     try:
       if not model.compiled_loss.built:
         model.compiled_loss.build(model.outputs)
       if not model.compiled_metrics.built:
+        # check if the loss = `sparse_categorical_crossentropy'
+        if ((isinstance(model.loss, str) and
+             ('sparse' in model.loss) and
+             (model.metrics_names[1].lower() in ['accuracy', 'acc']))):
+          logging.warning(
+              'The compiled metrics are yet to be built. Saving a model before '
+              'training or evaluating, may result in erroneous performance. '
+              'One workaround is to recompile the model before training.')
+        # check if the loss = tf.keras.losses.SparseCategoricalCrossentropy()
+        elif ((isinstance(model.loss.name, str)) and
+              ('sparse' in model.loss.name) and
+              (model.metrics_names[1].lower() in ['accuracy', 'acc'])):
+          logging.warning(
+              'The compiled metrics are yet to be built. Saving a model before '
+              'training or evaluating, may result in erroneous performance. '
+              'One workaround is to recompile the model before training.')
         model.compiled_metrics.build(model.outputs, model.outputs)
     except:  # pylint: disable=bare-except
       logging.warning(
@@ -332,5 +349,6 @@ def try_build_compiled_arguments(model):
 
 
 def is_hdf5_filepath(filepath):
+  """Checks whether the filepath includes .h5/.hdf5/.keras."""
   return (filepath.endswith('.h5') or filepath.endswith('.keras') or
           filepath.endswith('.hdf5'))