GAN training loop in PyTorch - ML Experiment: Train & Evaluate

Experiment - GAN training loop

Problem:You are training a Generative Adversarial Network (GAN) to generate images similar to a dataset of handwritten digits. The current training loop runs but the generator produces poor quality images and the discriminator quickly becomes too confident, causing unstable training.

Current Metrics:Generator loss: 1.5, Discriminator loss: 0.1, Generated image quality: low, Training stability: unstable

Issue:The GAN training loop lacks proper balance between generator and discriminator updates, causing the discriminator to overpower the generator early, leading to poor generator learning and unstable training.

Your Task

Improve the GAN training loop to stabilize training and improve generated image quality. Target: reduce discriminator loss to around 0.5, generator loss to below 1.0, and achieve stable training over 50 epochs.

Keep the model architectures unchanged.

Only modify the training loop and optimizer steps.

Use PyTorch framework.

Hint 1

Hint 2

Hint 3

Hint 4

Hint 5

Solution

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Simple GAN architecture placeholders
class Generator(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(100, 256),
            nn.ReLU(True),
            nn.Linear(256, 784),
            nn.Tanh()
        )
    def forward(self, z):
        return self.net(z)

class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(784, 256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )
    def forward(self, x):
        return self.net(x)

# Setup
batch_size = 64
lr = 0.0002
epochs = 50

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

generator = Generator().to(device)
discriminator = Discriminator().to(device)

criterion = nn.BCELoss()
optimizer_G = optim.Adam(generator.parameters(), lr=lr, betas=(0.5, 0.999))
optimizer_D = optim.Adam(discriminator.parameters(), lr=lr, betas=(0.5, 0.999))

for epoch in range(epochs):
    g_loss_epoch = 0.0
    d_loss_epoch = 0.0
    for real_imgs, _ in dataloader:
        real_imgs = real_imgs.view(-1, 784).to(device)
        batch_size_curr = real_imgs.size(0)

        # Labels
        real_labels = torch.ones(batch_size_curr, 1, device=device) * 0.9  # label smoothing
        fake_labels = torch.zeros(batch_size_curr, 1, device=device)

        # Train Discriminator
        optimizer_D.zero_grad()
        # Real images
        output_real = discriminator(real_imgs)
        loss_real = criterion(output_real, real_labels)
        # Fake images
        noise = torch.randn(batch_size_curr, 100, device=device)
        fake_imgs = generator(noise)
        output_fake = discriminator(fake_imgs.detach())  # detach to avoid generator gradients
        loss_fake = criterion(output_fake, fake_labels)

        d_loss = loss_real + loss_fake
        d_loss.backward()
        optimizer_D.step()

        # Train Generator
        optimizer_G.zero_grad()
        output_fake_for_g = discriminator(fake_imgs)
        g_loss = criterion(output_fake_for_g, real_labels)  # try to fool discriminator
        g_loss.backward()
        optimizer_G.step()

        g_loss_epoch += g_loss.item()
        d_loss_epoch += d_loss.item()

    avg_g_loss = g_loss_epoch / len(dataloader)
    avg_d_loss = d_loss_epoch / len(dataloader)
    print(f"Epoch {epoch+1}/{epochs} - Generator Loss: {avg_g_loss:.4f}, Discriminator Loss: {avg_d_loss:.4f}")

Added label smoothing for real labels to 0.9 to prevent discriminator overconfidence.

Used separate zero_grad calls before discriminator and generator backward passes.

Detached generator output when feeding fake images to discriminator to avoid gradient flow to generator during discriminator update.

Balanced training by updating discriminator once and generator once per batch.

Used Adam optimizer with betas (0.5, 0.999) for stable GAN training.

Results Interpretation

Before: Generator loss: 1.5, Discriminator loss: 0.1, unstable training with poor image quality.

After: Generator loss: ~0.85, Discriminator loss: ~0.55, stable training and better image quality.

Balancing the training steps of generator and discriminator, using label smoothing, and careful gradient management stabilizes GAN training and improves generated results.

Bonus Experiment

Try adding noise to the real images before feeding them to the discriminator to further improve training stability.

💡 Hint

Add small Gaussian noise to real images each batch to make discriminator less confident and encourage generator learning.

Practice

1. In a GAN training loop, what is the main role of the discriminator?

easy

A. To calculate the learning rate

B. To generate new fake data samples

C. To distinguish between real and fake data

D. To optimize the generator's weights directly

Solution

1. Step 1: Understand the discriminator's purpose

   The discriminator's job is to tell apart real data from fake data created by the generator.

2. Step 2: Differentiate roles in GAN

   The generator creates fake data, while the discriminator judges it. It does not generate data or adjust learning rates.

3. Final Answer:

   To distinguish between real and fake data -> Option C

4. Quick Check:

   Discriminator = spot fake data [OK]

Hint: Discriminator judges data authenticity, generator creates data [OK]

Common Mistakes:

• Confusing discriminator with generator
• Thinking discriminator generates data
• Assuming discriminator sets learning rate

2. Which of the following is the correct way to zero the gradients of both generator and discriminator optimizers in PyTorch during a GAN training loop?

easy

A. optimizer_G.clear(); optimizer_D.clear()

B. optimizer_G.zero_grad(); optimizer_D.zero_grad()

C. optimizer_G.reset_grad(); optimizer_D.reset_grad()

D. optimizer_G.zero(); optimizer_D.zero()

Solution

1. Step 1: Recall PyTorch optimizer method

   In PyTorch, the method to clear gradients before backpropagation is zero_grad().

2. Step 2: Check each option's validity

   Only optimizer_G.zero_grad(); optimizer_D.zero_grad() uses zero_grad(), which is the correct method. Other options use non-existent methods.

3. Final Answer:

   optimizer_G.zero_grad(); optimizer_D.zero_grad() -> Option B

4. Quick Check:

   Clear gradients = zero_grad() [OK]

Hint: Use zero_grad() to clear gradients before backward pass [OK]

Common Mistakes:

• Using non-existent optimizer methods
• Forgetting to zero gradients before backward
• Confusing zero_grad() with optimizer step()

3. Given the following PyTorch GAN training snippet, what will be the printed output for print(loss_D.item(), loss_G.item()) after one iteration?

real_data = torch.ones(4)
fake_data = torch.zeros(4)
criterion = torch.nn.BCELoss()

# Discriminator output
output_real = torch.tensor([0.9, 0.8, 0.95, 0.85])
output_fake = torch.tensor([0.1, 0.2, 0.05, 0.15])

loss_real = criterion(output_real, real_data)
loss_fake = criterion(output_fake, fake_data)
loss_D = (loss_real + loss_fake) / 2

# Generator tries to fool discriminator
output_fake_for_G = torch.tensor([0.4, 0.3, 0.5, 0.6])
loss_G = criterion(output_fake_for_G, real_data)

print(loss_D.item(), loss_G.item())

medium

A. Approximately 0.693 and 0.105

B. Approximately 0.0 and 1.0

C. Approximately 0.5 and 0.5

D. Approximately 0.136 and 0.831

Solution

1. Step 1: Calculate discriminator loss

   Loss_real = BCELoss([0.9,0.8,0.95,0.85], [1,1,1,1]) ≈ 0.136; Loss_fake = BCELoss([0.1,0.2,0.05,0.15], [0,0,0,0]) ≈ 0.136; Average loss_D ≈ 0.136.

2. Step 2: Calculate generator loss

   Generator tries to fool discriminator, so loss_G = BCELoss([0.4,0.3,0.5,0.6], [1,1,1,1]) ≈ 0.831.

3. Final Answer:

   Approximately 0.136 and 0.831 -> Option D

4. Quick Check:

   Discriminator loss < Generator loss [OK]

Hint: Discriminator loss low near 0, generator loss higher when fooled less [OK]

Common Mistakes:

• Mixing up real and fake labels in loss
• Confusing which loss belongs to generator
• Ignoring averaging of discriminator losses

4. Identify the error in this GAN training loop snippet and select the fix:

for real_data in dataloader:
    optimizer_D.zero_grad()
    output_real = discriminator(real_data)
    loss_real = criterion(output_real, torch.ones_like(output_real))

    fake_data = generator(torch.randn(batch_size, noise_dim))
    output_fake = discriminator(fake_data.detach())
    loss_fake = criterion(output_fake, torch.zeros_like(output_fake))

    loss_D = loss_real + loss_fake
    loss_D.backward()
    optimizer_D.step()

    optimizer_G.zero_grad()
    output_fake_for_G = discriminator(fake_data)
    loss_G = criterion(output_fake_for_G, torch.ones_like(output_fake_for_G))
    loss_G.backward()
    optimizer_D.step()

medium

A. Replace last optimizer_D.step() with optimizer_G.step()

B. Add optimizer_G.zero_grad() before loss_G.backward()

C. Detach fake_data before passing to generator

D. Use torch.zeros_like for output_fake_for_G labels

Solution

1. Step 1: Analyze optimizer steps

   After computing generator loss and backward, the generator's optimizer should update weights, not discriminator's.

2. Step 2: Identify incorrect optimizer call

   The code calls optimizer_D.step() twice; the second should be optimizer_G.step() to update generator.

3. Final Answer:

   Replace last optimizer_D.step() with optimizer_G.step() -> Option A

4. Quick Check:

   Generator update uses optimizer_G.step() [OK]

Hint: Generator updates need optimizer_G.step(), not optimizer_D.step() [OK]

Common Mistakes:

• Calling discriminator optimizer twice
• Forgetting to update generator weights
• Misplacing zero_grad calls

5. You want to improve GAN training stability by modifying the training loop. Which change below best helps prevent the discriminator from overpowering the generator?

hard

A. Train the discriminator less frequently than the generator

B. Increase the learning rate of the discriminator only

C. Remove the generator's loss calculation

D. Use the same optimizer instance for both models

Solution

1. Step 1: Understand discriminator overpowering

   If the discriminator learns too fast, the generator receives poor feedback and training stalls.

2. Step 2: Evaluate options to balance training

   Training discriminator less often than generator slows its learning, giving generator a chance to improve.

3. Final Answer:

   Train the discriminator less frequently than the generator -> Option A

4. Quick Check:

   Balance training frequency to stabilize GAN [OK]

Hint: Train discriminator less often to balance GAN training [OK]

Common Mistakes:

• Increasing discriminator learning rate worsens imbalance
• Removing generator loss stops generator learning
• Sharing optimizer causes parameter update conflicts