[bug] Training loop error in GANs

Enter the chapter number

Chapter-17: Autoencoders, GANs and Diffusion Models

Enter the page number

No response

What is the cell's number in the notebook

Cell 54

Enter the environment you are using to run the notebook

Kaggle

Describe your issue

Running the training loop gives

UserWarning: The model does not have any trainable weights. warnings.warn("The model does not have any trainable weights.")

error, the error is due to the updated Keras API and the discriminator and generator needs to be updated using tf.GradientTape() to override the train_step

Enter what you expected to happen

No response

If you found a workaround, describe it here

Creating a custom GAN class would be better here,

class GAN(keras.Model):
    def __init__(self, discriminator, generator, codings_size):
        super().__init__()
        self.discriminator = discriminator
        self.generator = generator
        self.codings_size = codings_size

        self.d_loss_tracker = tf.keras.metrics.Mean(name="d_loss")
        self.g_loss_tracker = tf.keras.metrics.Mean(name="g_loss")

       @property
       def metrics(self):
           return [self.d_loss_tracker,  self.g_loss_tracker]

      def compile(self, d_optimizer, g_optimizer, loss_fn):
         super().compile()
         self.d_optimizer = d_optimizer
         self.g_optimizer = g_optimizer
         self.loss_fn = loss_fn

     def train_step(self, real_images):
         batch_size = tf.shape(real_images)[0]

        # Train the discriminator

        random_latent_vectors = tf.random.normal(shape=(batch_size, self.codings_size))
        generated_images = self.generator(random_latent_vectors, training=True)
        combined_images = tf.concat([generated_images, real_images], axis=0)

        labels = tf.concat([tf.zeros((batch_size, 1)), tf.ones((batch_size, 1))], axis=0)

        with tf.GradientTape() as tape:
            predictions = self.discriminator(combined_images, training=True)
            d_loss = self.loss_fn(labels, predictions)

        d_gradients = tape.gradient(d_loss, self.discriminator.trainable_weights)
        self.d_optimizer.apply_gradients(zip(d_gradients, self.discriminator.trainable_weights))

        # Train the generator

        random_latent_vectors = tf.random.normal(shape=(batch_size, self.codings_size))
        misleading_labels = tf.ones((batch_size, 1))

        with tf.GradientTape() as tape:
            generated_images = self.generator(random_latent_vectors, training=True)
            predictions = self.discriminator(generated_images, training=True)
            g_loss = self.loss_fn(misleading_labels, predictions)

        g_gradients = tape.gradient(g_loss, self.generator.trainable_weights)
        self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_weights))

        self.d_loss_tracker.update_state(d_loss)
        self.g_loss_tracker.update_state(g_loss)

        return {"d_loss": self.d_loss_tracker.result(), "g_loss": self.g_loss_tracker.result()}