Generative Adversarial Network for MNIST: PyTorch

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.

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
import random
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from torchvision import datasets

print(f"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})
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")  # Deterministic cuBLAS on GPUs.
torch.use_deterministic_algorithms(True, warn_only=True)
DEVICE = torch.device(
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)
data_root = Path(os.getenv("HELIO_DATA_DIR", Path.home() / ".cache" / "helio-data-methods"))
mnist = datasets.MNIST(data_root, train=True, download=True)
x_train = (mnist.data.numpy().astype("float32") - 127.5) / 127.5
x_train = x_train[:, np.newaxis, ...]
loader = DataLoader(
    TensorDataset(torch.from_numpy(x_train)),
    batch_size=BATCH_SIZE,
    shuffle=True,
    drop_last=True,
    generator=torch.Generator().manual_seed(SEED),
)
fixed_noise = torch.randn(16, LATENT_DIM, 1, 1, generator=torch.Generator().manual_seed(SEED))

PyTorch 2.13.0+cu130

Define the generator and discriminator

The generator upsamples the noise from 1×1 to 7×7, 14×14, and 28×28 with transposed convolutions, with batch normalization between them and a tanh output. The discriminator mirrors it with two strided convolutions and dropout, and ends in one logit. Both use Adam with a small learning rate, as GAN training is unstable at larger steps.

generator = nn.Sequential(
    nn.ConvTranspose2d(LATENT_DIM, 256, 7, 1, 0, bias=False),
    nn.BatchNorm2d(256), nn.LeakyReLU(0.2),
    nn.ConvTranspose2d(256, 128, 5, 1, 2, bias=False),
    nn.BatchNorm2d(128), nn.LeakyReLU(0.2),
    nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False),
    nn.BatchNorm2d(64), nn.LeakyReLU(0.2),
    nn.ConvTranspose2d(64, 1, 4, 2, 1, bias=False), nn.Tanh(),
).to(DEVICE)
discriminator = nn.Sequential(
    nn.Conv2d(1, 64, 5, 2, 2), nn.LeakyReLU(0.2), nn.Dropout(0.3),
    nn.Conv2d(64, 128, 5, 2, 2), nn.LeakyReLU(0.2), nn.Dropout(0.3),
    nn.Flatten(), nn.Linear(128 * 7 * 7, 1),
).to(DEVICE)
criterion = nn.BCEWithLogitsLoss()
generator_optimizer = torch.optim.Adam(generator.parameters(), lr=1e-4)
discriminator_optimizer = torch.optim.Adam(discriminator.parameters(), lr=1e-4)

Train

Each batch performs one discriminator update (real images labelled 1, generated images 0) and one generator update (rewarding generated images the discriminator calls real). Images from the fixed noise are stored at a few epochs.

generator_losses, discriminator_losses, snapshots = [], [], {}
for epoch in range(EPOCHS):
    generator.train()
    epoch_generator, epoch_discriminator = [], []
    for (real_images,) in loader:
        real_images = real_images.to(DEVICE)
        noise = torch.randn(len(real_images), LATENT_DIM, 1, 1, device=DEVICE)
        generated_images = generator(noise)

        discriminator_optimizer.zero_grad()
        real_logits = discriminator(real_images)
        generated_logits = discriminator(generated_images.detach())
        real_targets = torch.ones_like(real_logits) * (1.0 - LABEL_SMOOTHING)
        discriminator_loss = criterion(real_logits, real_targets) + criterion(
            generated_logits, torch.zeros_like(generated_logits)
        )
        discriminator_loss.backward()
        discriminator_optimizer.step()

        generator_optimizer.zero_grad()
        generated_logits = discriminator(generated_images)
        generator_loss = criterion(generated_logits, torch.ones_like(generated_logits))
        generator_loss.backward()
        generator_optimizer.step()
        epoch_generator.append(generator_loss.item())
        epoch_discriminator.append(discriminator_loss.item())
    generator_losses.append(float(np.mean(epoch_generator)))
    discriminator_losses.append(float(np.mean(epoch_discriminator)))
    if epoch + 1 in SNAPSHOT_EPOCHS:
        generator.eval()
        with torch.no_grad():
            images = generator(fixed_noise.to(DEVICE)).cpu().numpy()
        snapshots[epoch + 1] = np.transpose(images, (0, 2, 3, 1))
        print(
            f"epoch {epoch + 1}: generator loss {generator_losses[-1]:.3f}, "
            f"discriminator loss {discriminator_losses[-1]:.3f}"
        )
generated = snapshots[EPOCHS]

epoch 1: generator loss 1.919, discriminator loss 0.678

epoch 5: generator loss 1.111, discriminator loss 1.139

epoch 25: generator loss 0.865, discriminator loss 1.284

epoch 50: generator loss 0.844, discriminator loss 1.305

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.294

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

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.