Use GANs for Image Generation

You understand how classifiers work, and maybe you've fine-tuned a model to label images. But here's the frustration: labeling is easy, creating is hard. How do you generate a brand-new, plausible image of a product, a face, or a landscape that has never existed — and do it programmatically, in Python? This lesson removes the mystery around Generative Adversarial Networks (GANs) and shows you a clean, practical path to implementing image generation with PyTorch, starting from a simple dense network and moving toward a DCGAN you can actually run on a laptop GPU.

By the end of this lesson, you will have a working mental model of adversarial training, a running code sample that generates synthetic digits (trained on MNIST), and a clear sense of where GANs fit inside the broader applied AI toolkit — and what to study next.

The problem this lesson solves

Most tutorials treat GANs like black magic: they show you a diagram, wave at a loss function, and then vanish. That leaves you with three real problems:

• You cannot debug a model you don't understand. If the generated images are nonsense, or the generator collapses to producing one boring image, you have no idea which component failed.
• You waste GPU hours. Without a clear step-by-step sequence, you will fiddle with hyperparameters blindly and burn through your quota.
• You cannot reuse the architecture. Every image domain (faces, products, medical scans) requires a slightly different generator and discriminator — but the training loop is always the same. If you only learn the "big name" GAN, you miss the transferable core.

This lesson gives you exactly that transferable core. By the end, you'll be able to take any GAN codebase, inspect the generator and discriminator, and immediately understand what each side is trying to do — and why.

Core concept / mental model

Think of a GAN as a counterfeiter and a detective playing an endless game.

• The generator is the counterfeiter. It receives a random noise vector (a "seed" of randomness) and tries to produce an image that looks real.
• The discriminator is the detective. It receives an image — either a real training image or a fake one from the generator — and tries to classify it as real or fake.

They train together, competing. The counterfeiter improves by learning what tricks fool the detective; the detective improves by spotting those tricks. After enough rounds, the counterfeiter becomes so good that the detective cannot tell the difference. That's the moment you have a useful generator.

Formally, a GAN optimizes a min-max game. The discriminator tries to maximize its ability to classify correctly; the generator tries to minimize the discriminator's accuracy. The loss curves look like they're fighting — that's expected, not a bug.

Key vocabulary to internalize:

• Latent space: the low-dimensional random input to the generator (often 100 or 128 numbers). Different points in this space map to different image styles.
• Adversarial loss: the combined objective for both networks.
• Mode collapse: when the generator finds one "safe" output that fools the discriminator and stops producing variety.
• DCGAN: a deep convolutional GAN — the standard architecture for images, using ConvTranspose2d instead of dense layers.

Pro tip: Never think of a GAN as a single model. It is two models fighting. Every time you see "GAN training," your mental image should be two curves tugging against each other, not a loss going down smoothly.

How it works step by step

Here is the complete training dance, step by step:

1. Prepare the real data. Load your image dataset (e.g., MNIST, CIFAR-10). Normalize images to [-1, 1] — this matches the generator's output activation (tanh).
2. Create the generator. It takes a noise vector z of shape (batch_size, latent_dim) and outputs an image tensor of shape (batch_size, channels, height, width).
3. Create the discriminator. It takes an image tensor and outputs a single logit (real vs. fake).
4. Define the loss. Use Binary Cross-Entropy (BCE) with logits. The discriminator's target is 1 for real and 0 for fake; the generator's target is 1 for its fakes (it wants to fool the discriminator).
5. Train the discriminator first, one batch at a time: separate real and fake batches, compute losses, backpropagate, and update the discriminator's weights.
6. Train the generator: use the discriminator's output on fake images, compute the loss against a target of 1, and backpropagate — but only update the generator's weights.
7. Repeat for many epochs. Periodically log the losses and generate a few sample images to see improvement.

The crucial trick is the order and separation: you must update the discriminator and generator in separate steps. When you train the generator, you never update the discriminator's weights.

Cause → effect: if the discriminator gets too good, the generator's gradient vanishes (loss flattens, learning stops). If the generator gets too good, the discriminator's loss shrinks to near zero and it no longer provides useful feedback. Balancing these two is the heart of GAN engineering.

Hands-on walkthrough

We'll build a DCGAN in PyTorch, trained on MNIST (28×28 grayscale digits). This is a complete, runnable example — you can paste it into a notebook or script.

First, install and import the essentials:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import torchvision.utils as vutils
import matplotlib.pyplot as plt

# Reproducibility
manualSeed = 42
torch.manual_seed(manualSeed)
print(f"PyTorch version: {torch.__version__}")

Now define the generator and discriminator as nn.Module subclasses:

class Generator(nn.Module):
    def __init__(self, latent_dim=100, ngf=64):
        super().__init__()
        self.main = nn.Sequential(
            # input: (batch, latent_dim, 1, 1)
            nn.ConvTranspose2d(latent_dim, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # state size: (ngf*8, 4, 4)
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # state size: (ngf*4, 8, 8)
            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # state size: (ngf*2, 16, 16)
            nn.ConvTranspose2d(ngf * 2, 1, 4, 2, 1, bias=False),
            nn.Tanh()
        )

    def forward(self, z):
        return self.main(z)

class Discriminator(nn.Module):
    def __init__(self, ndf=64):
        super().__init__()
        self.main = nn.Sequential(
            # input: (batch, 1, 28, 28)
            nn.Conv2d(1, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf * 2, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )

    def forward(self, img):
        return self.main(img)

Notice: the generator ends with a Tanh activation (outputs in [-1, 1]), and the discriminator ends with Sigmoid (outputs a probability). This matches the BCE loss.

Now set up the training loop. Here's a minimal version that trains for 5 epochs on MNIST:

def train_gan(dataloader, device, latent_dim=100, epochs=5):
    netG = Generator(latent_dim).to(device)
    netD = Discriminator().to(device)

    criterion = nn.BCELoss()
    lr = 0.0002
    beta1 = 0.5
    optG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))
    optD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))

    fixed_noise = torch.randn(64, latent_dim, 1, 1, device=device)
    print("Starting GAN training...")
    for epoch in range(epochs):
        for i, (real_imgs, _) in enumerate(dataloader):
            batch_size = real_imgs.size(0)
            real_imgs = real_imgs.to(device)
            real_labels = torch.ones(batch_size, 1, device=device)
            fake_labels = torch.zeros(batch_size, 1, device=device)

            # ---- Train discriminator ----
            netD.zero_grad()
            real_out = netD(real_imgs)
            lossD_real = criterion(real_out, real_labels)

            z = torch.randn(batch_size, latent_dim, 1, 1, device=device)
            fake_imgs = netG(z)
            fake_out = netD(fake_imgs.detach())
            lossD_fake = criterion(fake_out, fake_labels)

            lossD = lossD_real + lossD_fake
            lossD.backward()
            optD.step()

            # ---- Train generator ----
            netG.zero_grad()
            fake_out = netD(fake_imgs)  # detach not needed here
            lossG = criterion(fake_out, real_labels)
            lossG.backward()
            optG.step()

        # Log progress
        print(f"Epoch [{epoch+1}/{epochs}] LossD: {lossD.item():.4f}, LossG: {lossG.item():.4f}")
        with torch.no_grad():
            fake_fixed = netG(fixed_noise).detach().cpu()
            grid = vutils.make_grid(fake_fixed, normalize=True)
            plt.imshow(grid.permute(1, 2, 0))
            plt.title(f"Epoch {epoch+1}")
            plt.show()

# Load MNIST, normalize to [-1, 1]
dataset = datasets.MNIST(
    root='./data', train=True, download=True,
    transform=transforms.Compose([
        transforms.Resize(28),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5])
    ])
dataloader = DataLoader(dataset, batch_size=128, shuffle=True)

train_gan(dataloader, device='cuda' if torch.cuda.is_available() else 'cpu', epochs=5)

After 5 epochs, the output images should look like rough, blurry digits. After 30–50 epochs (with patience), they become sharp and diverse. The printed loss values will oscillate — that's normal. Expect LossD to hover around 0.8–1.2 and LossG around 0.8–1.5; if either drops to near zero, the other is failing.

Pro tip: Always fix a fixed_noise vector before training. This lets you see the same samples evolve each epoch, which is the clearest way to verify the generator is improving.

Compare options / when to choose what

GANs are not the only game in town. Here's how they compare to other generative approaches:

When to choose GANs:

• You need fast, one-shot generation (a single forward pass through the generator).
• You have a large dataset (GANs are data-hungry) and can afford a GPU.
• You want sharp, realistic images — GANs are better than VAEs at detail.

When to avoid GANs:

• You need stable, predictable training — diffusion or VAEs are safer.
• You have very little data (GANs will overfit and collapse).
• You need exact likelihood estimation — GANs don't give you probabilities.

Variations worth knowing

• Conditional GAN (cGAN): feed a class label to both networks to control the output (e.g., generate a "7" or a "9"). This is the foundation of many real applications.
• Wasserstein GAN (WGAN): replaces BCE loss with Earth-Mover distance for smoother training curves and fewer collapse issues — a common practical upgrade.
• StyleGAN: adds a mapping network and adaptive instance normalization for fine-grained control over image style — the architecture behind high-quality face synthesis.

Troubleshooting & edge cases

Loss goes to zero and stays there

• Symptom: LossD drops to 0.000 while LossG blows up.
• Cause: The discriminator has become too strong. It perfectly separates real from fake, providing no gradient signal.
• Fix: Lower the discriminator's learning rate, increase the generator's learning rate, or train the generator more often (e.g., 2–5 steps per discriminator step).

Mode collapse (generator outputs repetitive images)

• Symptom: All generated images are nearly identical, even with different noise vectors.
• Cause: The generator found a single output that fools the discriminator and has no incentive to vary.
• Fix: Use a WGAN loss (replace BCE with -mean(discriminator output)), add dropout to the discriminator, or increase the latent dimension.

Training is slow / OOM on GPU

• Symptom: CUDA out of memory or the notebook freezes.
• Cause: Batch size too large, or image resolution too high for your GPU.
• Fix: Reduce the batch size (e.g., 64), lower the image size (e.g., 28 for MNIST, 64 for CIFAR), or use gradient accumulation.

Generated images are all blurry

• Symptom: Outputs look like faded copies of each other.
• Cause: The generator is not training fast enough, or the discriminator is too weak.
• Fix: Train for more epochs (GANs need many). Increase the generator's capacity (ngf), or reduce the learning rate slightly for stability.

Values of images are weird (not in expected range)

• Symptom: Outputs are all black or all white.
• Cause: You prepared real data in [0,1] but the generator outputs [-1,1] (or vice versa).
• Fix: Always normalize real images to [-1,1] using transforms.Normalize([0.5], [0.5]) and keep Tanh at the generator output.

Blockquote pro tip: When debugging a GAN, log the real and fake loss separately. Start with a quick overfit test on a single batch: if the generator cannot memorize even one image, your architecture or optimizer is wrong.

What you learned & what's next

You now have a working mental model of GANs, a complete PyTorch DCGAN implementation for image generation, and a solid troubleshooting toolkit. Specifically, you can:

• Explain the adversarial training loop: generator vs. discriminator, separated updates, BCE loss.
• Implement a DCGAN for MNIST in PyTorch, from data loading to sampling.
• Debug common failure modes like mode collapse and vanishing gradients.
• Choose wisely between GANs, VAEs, and diffusion models based on your project's constraints.

This is just the beginning of generative AI in your Applied AI engineering path. The next logical step is handling larger, more varied datasets — such as CIFAR-10 color images — which forces you to scale your GAN architecture and training tricks. After that, you'll learn conditional generation, where you can ask the GAN for a specific class of image, a skill that directly powers applications like custom product mockups and data augmentation for underrepresented classes.

Keep the fixed_noise trick in your back pocket, and always watch those loss curves — they are the heartbeat of your GAN. Happy generating!