Generative Adversarial Network for MNIST: Keras 3

A generative adversarial network (GAN) trains two networks against each other. The generator turns a vector of 100 random numbers into a 28×28 image; the discriminator sees real and generated images and outputs the log-odds that its input is real. The generator improves by fooling the discriminator, the discriminator by not being fooled. This is the DCGAN design of Radford et al. (2016), trained on the 60,000 MNIST training images for 50 epochs with batches of 128. Reduce EPOCHS for a quicker run.

Keras and PyTorch

This notebook repeats the PyTorch version with the Keras 3 API, running on the same PyTorch backend. compile() sets the optimizer and loss, fit() runs the loop over batches and epochs, and callbacks such as early stopping replace hand-written control flow. The data split, model, and evaluation match the PyTorch notebook.

Setup and data

Pixel values are scaled to [-1, 1] to match the generator’s tanh output. A fixed set of 16 noise vectors is kept aside so the same samples can be followed through training.

import json
import os

os.environ["KERAS_BACKEND"] = "torch"
os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")  # Deterministic cuBLAS on GPUs.

import keras
import matplotlib.pyplot as plt
import numpy as np
import torch
from keras import layers

assert keras.backend.backend() == "torch"
print(f"Keras {keras.__version__}, backend {keras.backend.backend()}, PyTorch {torch.__version__}")

SEED = 42
EPOCHS = 50  # Reduce to 1 or 2 for a quicker run.
BATCH_SIZE = 128
LATENT_DIM = 100
LABEL_SMOOTHING = 0.0
SNAPSHOT_EPOCHS = sorted({1, max(1, EPOCHS // 10), max(1, EPOCHS // 2), EPOCHS})
keras.utils.set_random_seed(SEED)
torch.use_deterministic_algorithms(True, warn_only=True)
(x_train, _), _ = keras.datasets.mnist.load_data()
x_train = (x_train.astype("float32") - 127.5) / 127.5
x_train = x_train[..., np.newaxis]
dataset = torch.utils.data.DataLoader(
    torch.utils.data.TensorDataset(torch.from_numpy(x_train)),
    batch_size=BATCH_SIZE,
    shuffle=True,
    drop_last=True,
    generator=torch.Generator().manual_seed(SEED),
)
fixed_noise = keras.random.normal((16, LATENT_DIM), seed=SEED)

Keras 3.15.1, backend torch, PyTorch 2.13.0+cu130

Define the generator and discriminator

The generator projects the noise to 7×7×256 and upsamples twice with transposed convolutions to 28×28×1. The discriminator mirrors it with two strided convolutions and dropout. AdversarialModel overrides train_step so that fit() alternates one discriminator update and one generator update per batch.

generator = keras.Sequential(
    [
        keras.Input(shape=(LATENT_DIM,)),
        layers.Dense(7 * 7 * 256, use_bias=False),
        layers.BatchNormalization(),
        layers.LeakyReLU(negative_slope=0.2),
        layers.Reshape((7, 7, 256)),
        layers.Conv2DTranspose(128, 5, padding="same", use_bias=False),
        layers.BatchNormalization(),
        layers.LeakyReLU(negative_slope=0.2),
        layers.Conv2DTranspose(64, 5, strides=2, padding="same", use_bias=False),
        layers.BatchNormalization(),
        layers.LeakyReLU(negative_slope=0.2),
        layers.Conv2DTranspose(1, 5, strides=2, padding="same", activation="tanh"),
    ],
    name="generator",
)
discriminator = keras.Sequential(
    [
        keras.Input(shape=(28, 28, 1)),
        layers.Conv2D(64, 5, strides=2, padding="same"),
        layers.LeakyReLU(negative_slope=0.2),
        layers.Dropout(0.3),
        layers.Conv2D(128, 5, strides=2, padding="same"),
        layers.LeakyReLU(negative_slope=0.2),
        layers.Dropout(0.3),
        layers.Flatten(),
        layers.Dense(1),
    ],
    name="discriminator",
)


class AdversarialModel(keras.Model):
    def __init__(self, generator, discriminator, latent_dim, label_smoothing):
        super().__init__()
        self.generator = generator
        self.discriminator = discriminator
        self.latent_dim = latent_dim
        self.label_smoothing = label_smoothing
        self.seed_generator = keras.random.SeedGenerator(SEED)
        self.generator_loss_tracker = keras.metrics.Mean(name="generator_loss")
        self.discriminator_loss_tracker = keras.metrics.Mean(name="discriminator_loss")
        self.built = True

    @property
    def metrics(self):
        return [self.generator_loss_tracker, self.discriminator_loss_tracker]

    def compile(self, generator_optimizer, discriminator_optimizer, loss_function):
        super().compile()
        self.generator_optimizer = generator_optimizer
        self.discriminator_optimizer = discriminator_optimizer
        self.loss_function = loss_function

    def train_step(self, real_images):
        if isinstance(real_images, (tuple, list)):
            real_images = real_images[0]
        batch_size = real_images.shape[0]

        # Discriminator step: real images are labelled 1, generated images 0.
        noise = keras.random.normal((batch_size, self.latent_dim), seed=self.seed_generator)
        generated_images = self.generator(noise, training=True)
        self.zero_grad()
        real_logits = self.discriminator(real_images, training=True)
        generated_logits = self.discriminator(generated_images.detach(), training=True)
        real_targets = torch.ones_like(real_logits) * (1.0 - self.label_smoothing)
        discriminator_loss = self.loss_function(real_targets, real_logits) + self.loss_function(
            torch.zeros_like(generated_logits), generated_logits
        )
        discriminator_loss.backward()
        discriminator_weights = list(self.discriminator.trainable_weights)
        with torch.no_grad():
            self.discriminator_optimizer.apply(
                [weight.value.grad for weight in discriminator_weights], discriminator_weights
            )

        # Generator step: reward generated images that the discriminator calls real.
        noise = keras.random.normal((batch_size, self.latent_dim), seed=self.seed_generator)
        self.zero_grad()
        generated_logits = self.discriminator(self.generator(noise, training=True), training=True)
        generator_loss = self.loss_function(torch.ones_like(generated_logits), generated_logits)
        generator_loss.backward()
        generator_weights = list(self.generator.trainable_weights)
        with torch.no_grad():
            self.generator_optimizer.apply(
                [weight.value.grad for weight in generator_weights], generator_weights
            )

        self.generator_loss_tracker.update_state(generator_loss)
        self.discriminator_loss_tracker.update_state(discriminator_loss)
        return {
            "generator_loss": self.generator_loss_tracker.result(),
            "discriminator_loss": self.discriminator_loss_tracker.result(),
        }


gan = AdversarialModel(generator, discriminator, LATENT_DIM, LABEL_SMOOTHING)
gan.compile(
    generator_optimizer=keras.optimizers.Adam(1e-4),
    discriminator_optimizer=keras.optimizers.Adam(1e-4),
    loss_function=keras.losses.BinaryCrossentropy(from_logits=True),
)

Train

A callback stores the images generated from the fixed noise at a few epochs, so progress can be judged by eye.

snapshots = {}


class FixedNoiseSnapshots(keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        if epoch + 1 in SNAPSHOT_EPOCHS:
            images = generator(fixed_noise, training=False)
            snapshots[epoch + 1] = images.detach().cpu().numpy()


history = gan.fit(
    dataset, epochs=EPOCHS, verbose=0, shuffle=False, callbacks=[FixedNoiseSnapshots()]
)
generator_losses = [float(value) for value in history.history["generator_loss"]]
discriminator_losses = [float(value) for value in history.history["discriminator_loss"]]
generated = snapshots[EPOCHS]

Follow the fixed-noise samples

Each row shows the images generated from the same eight noise vectors at a later epoch. The loss curves follow; for a GAN they show the balance between the two networks, not the quality of the images.

fig, axes = plt.subplots(len(snapshots), 8, figsize=(9, 1.25 * len(snapshots) + 0.4), squeeze=False)
for row, (epoch, images) in zip(axes, sorted(snapshots.items())):
    for axis, image in zip(row, images):
        axis.imshow(image.squeeze(), cmap="gray", vmin=-1, vmax=1)
        axis.set_xticks([])
        axis.set_yticks([])
    row[0].set_ylabel(f"epoch {epoch}", rotation=0, ha="right", va="center")
fig.suptitle("Images generated from fixed noise")
plt.tight_layout()
plt.show()

fig, ax = plt.subplots(figsize=(7, 3.4))
ax.plot(np.arange(1, len(generator_losses) + 1), generator_losses, label="generator")
ax.plot(np.arange(1, len(discriminator_losses) + 1), discriminator_losses, label="discriminator")
ax.axhline(np.log(2), color="0.6", linestyle=":", label="generator loss at equilibrium, ln 2")
ax.set(title="Adversarial losses", xlabel="Epoch", ylabel="Binary cross-entropy")
ax.legend()
ax.grid(alpha=0.25)
plt.show()

pixel_diversity = float(generated.std(axis=0).mean())
print(f"mean per-pixel standard deviation across the 16 final samples: {pixel_diversity:.3f}")

mean per-pixel standard deviation across the 16 final samples: 0.316

What the results show

The samples sharpen from blobs into recognizable digits within the first few epochs and improve slowly after that. The losses settle near a fixed ratio rather than falling, which is the expected behaviour when neither network dominates. A GAN has no held-out score as simple as test accuracy: the per-pixel spread across samples only guards against collapse onto a single image, and judging whether the generator covers all ten digits in the right proportions needs a separate classifier or a metric such as the Fréchet inception distance.

Try it yourself in Keras

Change one choice at a time and keep the data split and evaluation unchanged:

• set LABEL_SMOOTHING to 0.1 and compare the samples and losses.

• halve the discriminator’s learning rate and watch the balance of the two losses.

• reduce LATENT_DIM to 10 and look for less variety in the samples.