Skip to content
← Generative AI
Q71

Basic GAN for generating handwritten digits

GANMNISTAdversarial Training
Implement a basic GAN using Python and TensorFlow or PyTorch on the MNIST dataset. Include data preprocessing, Generator and Discriminator construction, adversarial training, loss functions, and visualization of synthetic handwritten digits at different training stages.

What's actually being asked

Build and train a vanilla GAN on MNIST in PyTorch: preprocessing, a Generator, a Discriminator, the adversarial training loop with correct loss signs and .detach() usage, and a visualization of generated digits at multiple training stages (not just the final epoch).

Answer

The adversarial setup

G maps noise z ~ N(0,I) to a fake image G(z); D maps an image to a probability that it's real. They're trained against each other: min_G max_D V(D,G) = 𝔼_{x~p_data}[log D(x)] + 𝔼_{z~p_z}[log(1 − D(G(z)))]. D wants to maximize this; G wants to minimize it — but in practice G is trained with the non-saturating loss (maximize log D(G(z))) instead of the original saturating one, since it gives much stronger gradients early on when D easily rejects G's (still-bad) samples.

Full code

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.utils import make_grid
import matplotlib.pyplot as plt
 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
# --- Preprocessing: scale to [-1, 1] to match the Generator's Tanh output ---
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,)),
])
train_data = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
 
latent_dim = 100
 
class Generator(nn.Module):
    def __init__(self, latent_dim=100):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(latent_dim, 256), nn.LeakyReLU(0.2),
            nn.Linear(256, 512), nn.LeakyReLU(0.2),
            nn.Linear(512, 1024), nn.LeakyReLU(0.2),
            nn.Linear(1024, 784), nn.Tanh(),
        )
 
    def forward(self, z):
        return self.net(z).view(-1, 1, 28, 28)
 
class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(784, 512), nn.LeakyReLU(0.2),
            nn.Linear(512, 256), nn.LeakyReLU(0.2),
            nn.Linear(256, 1), nn.Sigmoid(),
        )
 
    def forward(self, x):
        return self.net(x.view(x.size(0), -1))
 
G = Generator(latent_dim).to(device)
D = Discriminator().to(device)
 
criterion = nn.BCELoss()
opt_G = torch.optim.Adam(G.parameters(), lr=2e-4, betas=(0.5, 0.999))
opt_D = torch.optim.Adam(D.parameters(), lr=2e-4, betas=(0.5, 0.999))
 
fixed_noise = torch.randn(64, latent_dim).to(device)
snapshots = {}
 
epochs = 20
for epoch in range(epochs):
    for real_images, _ in train_loader:
        real_images = real_images.to(device)
        batch_size = real_images.size(0)
        real_labels = torch.ones(batch_size, 1).to(device)
        fake_labels = torch.zeros(batch_size, 1).to(device)
 
        # --- Train Discriminator: maximize log D(x) + log(1 - D(G(z))) ---
        z = torch.randn(batch_size, latent_dim).to(device)
        fake_images = G(z)
 
        d_real = D(real_images)
        d_fake = D(fake_images.detach())      # detach: don't backprop into G during D's step
        loss_D = criterion(d_real, real_labels) + criterion(d_fake, fake_labels)
 
        opt_D.zero_grad()
        loss_D.backward()
        opt_D.step()
 
        # --- Train Generator: non-saturating loss, maximize log D(G(z)) ---
        z = torch.randn(batch_size, latent_dim).to(device)
        fake_images = G(z)
        d_fake_for_g = D(fake_images)
        loss_G = criterion(d_fake_for_g, real_labels)   # label fakes "real" => minimizing this maximizes log D(G(z))
 
        opt_G.zero_grad()
        loss_G.backward()
        opt_G.step()
 
    print(f"Epoch {epoch+1}/{epochs}  Loss_D: {loss_D.item():.4f}  Loss_G: {loss_G.item():.4f}")
 
    if epoch + 1 in (1, 5, 10, 20):
        with torch.no_grad():
            samples = G(fixed_noise).cpu()
        snapshots[epoch + 1] = samples
 
# --- Visualize synthetic digits at different training stages ---
fig, axes = plt.subplots(1, len(snapshots), figsize=(4 * len(snapshots), 4))
for ax, (ep, samples) in zip(axes, snapshots.items()):
    grid = make_grid(samples, nrow=8, normalize=True)
    ax.imshow(grid.permute(1, 2, 0).squeeze(), cmap="gray")
    ax.set_title(f"Epoch {ep}")
    ax.axis("off")
plt.show()

Mapping code to theory

  • Preprocessing: Normalize((0.5,), (0.5,)) scales pixels from [0,1] to [-1,1], matching the Generator's final Tanh output range — a real/common GAN detail that trips people up if skipped (output range mismatch stalls training).
  • Generator: noise → wider and wider FC layers with LeakyReLU (keeps gradients alive even for negative inputs, important since D's gradient signal to G is already weak early on) → Tanh output reshaped to a 28×28 image.
  • Discriminator: a simple binary classifier, Sigmoid output = P(real).
  • .detach(): when training D, fake_images.detach() cuts the computation graph so D's backward pass doesn't touch G's parameters at all — G is updated only in its own, separate step, on a fresh forward pass through D.
  • Loss signs: BCELoss computes −[y log ŷ + (1−y) log(1−ŷ)]; feeding (d_real, 1) and (d_fake, 0) to it and summing gives exactly −log D(x) − log(1−D(G(z))), so minimizing loss_D is the same as maximizing log D(x) + log(1 − D(G(z))), matching the canonical formula. Feeding (d_fake_for_g, 1) for the generator step is the standard "trick" that implements the non-saturating maximize log D(G(z)) objective via the same BCELoss call.

Visualization across training stages

Snapshots are saved at epochs 1, 5, 10, and 20 (from the same fixed_noise, so the comparison isolates training progress rather than noise variation): expect epoch 1 to look like unstructured gray noise, epoch 5 to show vague blob-like digit shapes, and by epoch 10–20 recognizable (if still somewhat rough, for this simple FC architecture) handwritten digits.

Exam tip

If asked why D is trained with .detach()'d fakes but G's step re-runs D(fake_images) without detaching: G's step needs gradients to flow back through D into G (that's the only way G gets any training signal — it never sees real images directly), whereas D's step must not update G, hence the detach there specifically.