Repository: keras-team/keras
Issue category: Bug Report

Issue:
NO _loss_tracker on train_on_batch because compile model multiple times. Possible Bug. 
Same code of a GAN works perfectly in Keras 3.3 but doesn´t work in keras 3.6. There is an error on train_on_batch I believe is because an bug introduce in a change in Keras 3.6 
The code is this:

    import numpy as np
    import matplotlib.pyplot as plt
    import random
    
    from keras.datasets import mnist
    from keras.utils import plot_model
    from keras.models import Sequential, Model
    from keras.layers import (Input, Conv2D, Dense, Activation, 
                             Flatten, Reshape, Dropout,
                             UpSampling2D, MaxPooling2D,
                             BatchNormalization, LeakyReLU, Conv2DTranspose,
                             GlobalMaxPooling2D)
    from keras.losses import BinaryCrossentropy
    from keras.optimizers import Adam
    from keras.metrics import Mean, Accuracy
    from keras.backend import backend
    from keras.random import SeedGenerator, uniform, normal
    from keras import ops
    
    !pip list | grep keras
    
    (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
    X_train = X_train.astype('float32')/127.5 -1
    X_train = np.expand_dims(X_train, axis = 3)
    
    def create_generator():
        generator = Sequential(
            [
                Input(shape = (100,)),
                Dense(7 * 7 * 128),
                LeakyReLU(0.2),
                Reshape((7, 7, 128)),
                Conv2DTranspose(128, 4, 2, "same"),
                LeakyReLU(0.2),
                Conv2DTranspose(256, 4, 2, "same"),
                LeakyReLU(0.2),
                Conv2D(1, 7, padding = "same", activation = "tanh"),
            ],
            name = "generator",
        )
        return generator
    
    def create_discriminator():
        discriminator = Sequential(
            [
                Input(shape = (28, 28, 1)),
                Conv2D(64, 3, 2, "same"),
                LeakyReLU(0.2),
                Conv2D(128, 3, 2, "same"),
                LeakyReLU(0.2),
                Conv2D(256, 3, 2, "same"),
                LeakyReLU(0.2),
                Flatten(),
                Dropout(0.2),
                Dense(1, activation = "sigmoid"),
            ],
            name = "discriminator",
        )
        return discriminator
    
    generator = create_generator()
    discriminator = create_discriminator()
    
    discriminator.compile(loss = 'binary_crossentropy', optimizer = Adam(), metrics = ['accuracy'])
    discriminator.trainable = False
    
    
    ###Print for debugging/show the error 
    print('---Debugging/show the error after compiled discriminator and before compiled combined---')
    print('discriminator.compiled ->', discriminator.compiled)
    print('discriminator.optimizer ->', discriminator.optimizer)
    print('discriminator.train_function ->', discriminator.train_function)
    print('discriminator.train_step ->', discriminator.train_step)
    print('discriminator.metrics ->', discriminator.metrics)
    print('discriminator._loss_tracker ->', discriminator._loss_tracker)
    print('discriminator._jit_compile ->', discriminator._jit_compile)
    ###
    
    z = Input(shape=(100,))
    img = generator(z)
    validity = discriminator(img)
    
    combined = Model(z, validity)
    combined.compile(loss = 'binary_crossentropy', optimizer = Adam())
    
    ###Print for debugging/show the error
    print('---Debugging/show the error after compiled discriminator and combined---')
    print('discriminator.compiled ->', discriminator.compiled)
    print('discriminator.optimizer ->', discriminator.optimizer)
    print('discriminator.train_function ->', discriminator.train_function)
    print('discriminator.train_step ->', discriminator.train_step)
    print('discriminator.metrics ->', discriminator.metrics)
    print('discriminator._loss_tracker ->', discriminator._loss_tracker)
    print('discriminator._jit_compile ->', discriminator._jit_compile)
    ###
    
    def train(X_train, generator, discriminator, combined, epochs, batch_size = 32, sample_interval = 100):
        valid = np.ones((batch_size, 1))
        fake = np.zeros((batch_size, 1))
    
        history = {
            'd_loss' : [],
            'd_acc' : [],
            'g_loss' : []
        }
    
        for epoch in range(epochs):
            print("----EPOCH " + str(epoch) + '-----')
            for batch in range(int(len(X_train)/batch_size)):
                #  Train the Discriminator
                noise = np.random.normal(0, 1, (batch_size, 100))
                gen_imgs = generator.predict(noise, verbose = 0)
                imgs = X_train[batch*batch_size : (batch+1)*batch_size]
    
                #Print for debugging/show the error
                print('---Debugging/show the error---')
                print('discriminator.compiled ->', discriminator.compiled)
                print('discriminator.optimizer ->', discriminator.optimizer)
                print('discriminator._loss_tracker ->', discriminator._loss_tracker)
                print('discriminator._jit_compile ->', discriminator._jit_compile)
    
    
                d_loss_real = discriminator.train_on_batch(imgs, valid)
                d_loss_fake = discriminator.train_on_batch(gen_imgs, fake)
                d_loss = 0.5 * np.add(d_loss_real, d_loss_fake)
    
                # Train the Generator
                noise = np.random.normal(0, 1, (batch_size, 100))
                g_loss = combined.train_on_batch(noise, valid)
    
                # Save losses
                history['d_loss'].append(d_loss[0])
                history['d_acc'].append(d_loss[1])
                history['g_loss'].append(g_loss[0])
    
    train(X_train, generator, discriminator, combined, epochs = 2, batch_size = 256, sample_interval = 100)
    
    The error is this:
    
        Cell In[13], line 128, in train(X_train, generator, discriminator, combined, epochs, batch_size, sample_interval)
            124 print('discriminator._loss_tracker ->', discriminator._loss_tracker)
            125 print('discriminator._jit_compile ->', discriminator._jit_compile)
        --> 128 d_loss_real = discriminator.train_on_batch(imgs, valid)
            129 d_loss_fake = discriminator.train_on_batch(gen_imgs, fake)
            130 d_loss = 0.5 * np.add(d_loss_real, d_loss_fake)
        
        File ~/.local/lib/python3.10/site-packages/keras/src/backend/torch/trainer.py:468, in TorchTrainer.train_on_batch(self, x, y, sample_weight, class_weight, return_dict)
            465 self._symbolic_build(data_batch=data)
            466 self.make_train_function()
        --> 468 logs = self.train_function([data])
            469 logs = tree.map_structure(lambda x: np.array(x), logs)
            470 if return_dict:
        
        File ~/.local/lib/python3.10/site-packages/keras/src/backend/torch/trainer.py:117, in TorchTrainer.make_train_function.<locals>.one_step_on_data(data)
            115 """Runs a single training step on a batch of data."""
            116 data = data[0]
        --> 117 return self.train_step(data)
        
        File ~/.local/lib/python3.10/site-packages/keras/src/backend/torch/trainer.py:55, in TorchTrainer.train_step(self, data)
             50 self.zero_grad()
             52 loss = self._compute_loss(
             53     x=x, y=y, y_pred=y_pred, sample_weight=sample_weight, training=True
             54 )
        ---> 55 self._loss_tracker.update_state(
             56     loss, sample_weight=tree.flatten(x)[0].shape[0]
             57 )
             58 if self.optimizer is not None:
             59     loss = self.optimizer.scale_loss(loss)
        
        AttributeError: 'NoneType' object has no attribute 'update_state'.

The error says that self._loss_tracker.update_state is None when is should be metrics_module.Mean(name="loss") as is been compiled.
The print that I write in the code shows that after compiled the discriminator and before the combined:

    ---Debugging/show the error after compiled discriminator and before compiled combined---
    discriminator.compiled -> True
    discriminator.optimizer -> <keras.src.optimizers.adam.Adam object at 0x77ecf6bb0a00>
    discriminator.train_function -> None
    discriminator.train_step -> <bound method TensorFlowTrainer.train_step of <Sequential name=discriminator, built=True>>
    discriminator.metrics -> [<Mean name=loss>, <CompileMetrics name=compile_metrics>]
    discriminator._loss_tracker -> <Mean name=loss>
    discriminator._jit_compile -> True

However after compiled the combined:
  discriminator.compiled -> True
  discriminator.optimizer -> <keras.src.optimizers.adam.Adam object at 0x77ecf6bb0a00>
  discriminator.train_function -> None
  discriminator.train_step -> <bound method TensorFlowTrainer.train_step of <Sequential name=discriminator, built=True>>
  discriminator.metrics -> []
  discriminator._loss_tracker -> None
  discriminator._jit_compile -> True

So the problems realives that compiling the combined erases the metrics (loss_tracks,...) of the discriminator what it shouldn't and keeps the discriminator as compiled when it shouldn´t becouse it undo the compiled. I belive the bug relieves in a change introduce in Keras 3.6 that  change the compile of keras/src/trainers/trainer.py:
![aux](https://github.com/user-attachments/assets/02722eb9-c962-4244-9151-85373504e5ed)
The function self._clear_previous_trainer_metrics not only clears the metrics of combined but also of discriminator what that makes that discriminator not having proper metrics. 
![aux2](https://github.com/user-attachments/assets/356c42f1-f09b-4141-84f4-fe4ab99eb972).
My pull request to this possible error is: https://github.com/keras-team/keras/pull/20473
I try the code with the thee backeends and happens always
I hope it help ! :)



Hints:
any update on this bug ?
Hi @TheMGGdev and @mohammad-rababah -

Here getting the error because you are doing `discriminator.compile(loss = 'binary_crossentropy', optimizer = Adam(), metrics = ['accuracy'])`and then freeze the weights during training by `discriminator.trainable = False` . 

So when trying to compile combine model  `combined.compile(loss = 'binary_crossentropy', optimizer = Adam())`, discriminator weights are frozen due to  `traininable=False`. 

That's why `discriminator.train_function` will become `None` on update_state. 

You can compile discriminator after the combine model compile. It will resolve your error. 

```
z = Input(shape=(100,))
img = generator(z)
validity = discriminator(img)
combined = Model(z, validity)
combined.compile(loss = 'binary_crossentropy', optimizer = Adam())
discriminator.compile(loss = 'binary_crossentropy', optimizer = Adam(), metrics = ['accuracy'])
discriminator.trainable = False
```

Attached [gist](https://colab.sandbox.google.com/gist/mehtamansi29/71713567be35f3b6be46974fa7c63660/20474-no-_loss_tracker-on-train_on_batch-because-compile-model-multiple-times-possible-bug.ipynb) for your reference. 
The error has nothing to do with that. There are trainings in which a model that is a combination of several models, you don't want to train one of them, as in this example with the GANS. Here you have generator, discriminator and combined (which is generator + discriminator). When you create the combined model, which is the one you are going to use to train the generator, you want the discriminator not to train, so you put discriminator.trainable = False, Explained more simply:

- Discriminator training: discriminator.train_on_batch we generate with the discriminator and the discriminator learns.
- Generator training: combined.train_on_batch we generate with the combined (discriminator + generator)  the combined learns but as we want only the generator to learn we set before we create the combined that the discriminator doesn´t learn.

The code works perfectly in other versions and it comes from a example of this repo https://github.com/eriklindernoren/Keras-GAN/blob/master/cgan/cgan.py. The real problem is that the new versions of Keras when you do the combined, it clears all the trainer metrics and therefore when you do the combined  compile it deletes the discriminator metrics. As explained in the following pull request https://github.com/keras-team/keras/pull/20473 . Put the 

    discriminator.compile(loss = ‘binary_crossentropy’, optimizer = Adam(), metrics = [‘accuracy’]) '
    discriminator.trainable = False 

after the combined compile instead of before does not solve the problem but creates another one, because now in the generator training the generator and the discriminator will be trained.  Hope it helps :)
@TheMGGdev looks like we just need a unit test to go along with your change https://github.com/keras-team/keras/pull/20473

Patch:
diff --git a/keras/src/trainers/trainer.py b/keras/src/trainers/trainer.py
index ff103b535c3..27bcfea9381 100644
--- a/keras/src/trainers/trainer.py
+++ b/keras/src/trainers/trainer.py
@@ -140,7 +140,6 @@ def compile(
                 wrapped in a `LossScaleOptimizer`, which will dynamically
                 scale the loss to prevent underflow.
         """
-        self._clear_previous_trainer_metrics()
         optimizer = optimizers.get(optimizer)
         self.optimizer = optimizer
         if (
@@ -287,21 +286,6 @@ def _get_own_metrics(self):
         metrics.extend(self._metrics)
         return metrics
 
-    def _clear_previous_trainer_metrics(self):
-        for layer in self._flatten_layers(include_self=False):
-            if not isinstance(layer, Trainer):
-                continue
-            # A sublayer might be a Trainer. In that case, we need to clear
-            # the Trainer-related metrics, as they are not usable when a
-            # new Trainer is instantiated.
-            for m in self._get_own_metrics():
-                layer._tracker.untrack(m)
-            layer._loss_tracker = None
-            layer._compile_metrics = None
-            if layer._compile_loss is not None:
-                layer._compile_loss._metrics.clear()
-            layer._metrics.clear()
-
     def compute_loss(
         self,
         x=None,


Test patch:
diff --git a/keras/src/trainers/trainer_test.py b/keras/src/trainers/trainer_test.py
index 434fe47c969..4e25f4d3437 100644
--- a/keras/src/trainers/trainer_test.py
+++ b/keras/src/trainers/trainer_test.py
@@ -407,6 +407,38 @@ def test_nested_trainer_metrics_without_compile(self):
         self.assertEqual(new_model.metrics[0], new_model._loss_tracker)
         self.assertEqual(new_model.metrics[1], new_model._compile_metrics)
 
+    def test_multiple_compiles(self):
+        # https://github.com/keras-team/keras/issues/20474
+        model1 = ExampleModel(units=3)
+        model2 = ExampleModel(units=3)
+        model1.compile(
+            optimizer=optimizers.SGD(),
+            loss=losses.MeanSquaredError(),
+            metrics=[metrics.MeanSquaredError()],
+        )
+
+        # Combine these 2 models into `combined`.
+        inputs = keras.Input(shape=(4,))
+        x = model1(inputs)
+        outputs = model2(x)
+        combined = models.Model(inputs, outputs)
+        combined.compile(
+            optimizer=optimizers.SGD(),
+            loss=losses.MeanSquaredError(),
+            metrics=[metrics.MeanSquaredError()],
+        )
+
+        self.assertLen(model1.metrics, 2)
+        self.assertIsNotNone(model1._loss_tracker)
+        self.assertEqual(model1.metrics[0], model1._loss_tracker)
+        self.assertEqual(model1.metrics[1], model1._compile_metrics)
+
+        # `combined.metrics` will not include `model1.metrics`.
+        self.assertLen(combined.metrics, 2)
+        self.assertIsNotNone(combined._loss_tracker)
+        self.assertEqual(combined.metrics[0], combined._loss_tracker)
+        self.assertEqual(combined.metrics[1], combined._compile_metrics)
+
     @pytest.mark.skipif(
         backend.backend() != "torch",
         reason="torch backend runs in eager mode for jit_compile='auto'",