GAN Training Pattern

Why a GAN breaks the single-loss mold

A GAN is two networks locked in a minimax game. The discriminator learns to tell real images from generated ones; the generator learns to produce fakes the discriminator can't catch. They have opposing objectives, separate weights, and separate optimizers — so the single loss, single optimizer shape that fit() assumes simply can't express them. Each batch needs two distinct update steps in sequence.

Why train_step() still fits

It's tempting to reach straight for a manual loop here, but the alternating pattern lives perfectly inside a train_step() override. You hold both sub-models on the GAN object, and inside one step you train the discriminator on a real+fake batch, then train the generator against frozen-discriminator feedback, and return both losses. Because you stayed inside fit(), the whole GAN inherits callbacks, checkpointing, and distribution for free. The skeleton below keeps the two updates in named helper methods (_train_discriminator / _train_generator) so the step reads as the game it models; each helper holds its own gradient tape and optimizer.

Code

GAN as a train_step() override·python

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

    def train_step(self, real_images):
        batch_size = keras.ops.shape(real_images)[0]
        noise = keras.random.normal(
            shape=(batch_size, self.latent_dim)
        )

        # Train discriminator on real + fake
        fake_images = self.generator(noise)
        combined = keras.ops.concatenate([real_images, fake_images])
        labels = keras.ops.concatenate([
            keras.ops.ones((batch_size, 1)),
            keras.ops.zeros((batch_size, 1)),
        ])
        d_loss = self._train_discriminator(combined, labels)

        # Train generator to fool the discriminator
        noise = keras.random.normal(shape=(batch_size, self.latent_dim))
        misleading_labels = keras.ops.ones((batch_size, 1))
        g_loss = self._train_generator(noise, misleading_labels)

        return {"d_loss": d_loss, "g_loss": g_loss}

Why a GAN breaks the single-loss mold

Why train_step() still fits

Code

External links

Exercise

Progress

Comments 0