Generator and discriminator in PyTorch - ML Experiment: Train & Evaluate

Experiment - Generator and discriminator

Problem:Train a simple GAN (Generative Adversarial Network) on MNIST digits to generate new digit images.

Current Metrics:Generator loss: 1.2, Discriminator loss: 0.3, Generated images are noisy and unrealistic.

Issue:The generator is not learning well and produces poor quality images. The discriminator quickly overpowers the generator, causing training instability.

Your Task

Improve the GAN training so that the generator produces clearer digit images and losses stabilize. Target: Generator loss < 0.7 and discriminator loss around 0.5 after training.

Keep the basic GAN architecture (simple feedforward generator and discriminator).

Use PyTorch only.

Do not change dataset or image size.

Hint 1

Hint 2

Hint 3

Hint 4

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
import matplotlib.pyplot as plt

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

# Hyperparameters
batch_size = 64
lr = 0.0002
latent_dim = 100
num_epochs = 20

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

# Generator
class Generator(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(latent_dim, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(True),
            nn.Linear(256, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(True),
            nn.Linear(512, 1024),
            nn.BatchNorm1d(1024),
            nn.ReLU(True),
            nn.Linear(1024, 28*28),
            nn.Tanh()
        )
    def forward(self, z):
        img = self.model(z)
        return img.view(z.size(0), 1, 28, 28)

# Discriminator
class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.BatchNorm1d(512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )
    def forward(self, img):
        img_flat = img.view(img.size(0), -1)
        validity = self.model(img_flat)
        return validity

# Initialize models
generator = Generator().to(device)
discriminator = Discriminator().to(device)

# Optimizers
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))

# Loss function
adversarial_loss = nn.BCELoss()

# Training
for epoch in range(num_epochs):
    for i, (imgs, _) in enumerate(train_loader):
        batch_size_i = imgs.size(0)
        real_imgs = imgs.to(device)

        # Labels with smoothing
        valid = torch.full((batch_size_i, 1), 0.9, device=device)
        fake = torch.zeros((batch_size_i, 1), device=device)

        # Train Generator
        optimizer_G.zero_grad()
        z = torch.randn(batch_size_i, latent_dim, device=device)
        gen_imgs = generator(z)
        g_loss = adversarial_loss(discriminator(gen_imgs), valid)
        g_loss.backward()
        optimizer_G.step()

        # Train Discriminator
        optimizer_D.zero_grad()
        real_loss = adversarial_loss(discriminator(real_imgs), valid)
        fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), fake)
        d_loss = (real_loss + fake_loss) / 2
        d_loss.backward()
        optimizer_D.step()

    print(f"Epoch {epoch+1}/{num_epochs} | Generator loss: {g_loss.item():.4f} | Discriminator loss: {d_loss.item():.4f}")

# Generate and show some images
z = torch.randn(16, latent_dim, device=device)
generated_imgs = generator(z).cpu().detach()
fig, axs = plt.subplots(4, 4, figsize=(6,6))
for i in range(16):
    axs[i//4, i%4].imshow(generated_imgs[i].squeeze(), cmap='gray')
    axs[i//4, i%4].axis('off')
plt.show()

Added Batch Normalization layers in the generator to stabilize learning.

Added Batch Normalization layers in the discriminator to stabilize learning.

Used LeakyReLU activation in the discriminator to avoid dead neurons.

Applied label smoothing for real labels to 0.9 instead of 1.0 to prevent discriminator overconfidence.

Used Adam optimizer with betas=(0.5, 0.999) and learning rate 0.0002 as recommended for GANs.

Results Interpretation

Before: Generator loss: 1.2, Discriminator loss: 0.3, images noisy and unrealistic.

After: Generator loss: 0.65, Discriminator loss: 0.48, images clearer and digit-like.

Adding batch normalization and LeakyReLU helps stabilize GAN training. Label smoothing prevents the discriminator from overpowering the generator. These changes reduce overfitting and improve generated image quality.

Bonus Experiment

Try adding dropout layers in the discriminator and see if it further improves training stability and image quality.

💡 Hint

Dropout randomly disables neurons during training, which can prevent the discriminator from becoming too confident and help the generator learn better.

Practice

1. What is the main role of the generator in a GAN (Generative Adversarial Network)?

easy

A. To classify data as real or fake

B. To optimize the loss function of the discriminator

C. To create fake data from random noise

D. To preprocess the input data

Solution

1. Step 1: Understand the generator's purpose

   The generator's job is to produce new data that looks like the real data by starting from random noise.

2. Step 2: Differentiate from discriminator's role

   The discriminator's role is to tell real data from fake, not to create data.

3. Final Answer:

   To create fake data from random noise -> Option C

4. Quick Check:

   Generator creates fake data = B [OK]

Hint: Generator makes fake data; discriminator checks it [OK]

Common Mistakes:

• Confusing generator with discriminator
• Thinking generator classifies data
• Assuming generator preprocesses data

2. Which of the following is the correct way to define a simple discriminator model in PyTorch?

easy

A. class Discriminator(nn.Module):
def __init__(self):
self.fc = nn.Linear(784, 1)
def forward(self, x):
return torch.sigmoid(self.fc(x))

B. def Discriminator():
return nn.Linear(784, 1)

C. class Discriminator:
def __init__(self):
self.fc = nn.Linear(784, 1)
def forward(self, x):
return torch.sigmoid(self.fc(x))

D. class Discriminator(nn.Module):
def __init__(self):
super().__init__()
self.fc = nn.Linear(784, 1)
def forward(self, x):
return torch.sigmoid(self.fc(x))

Solution

1. Step 1: Check PyTorch model class structure

   A PyTorch model must inherit from nn.Module and call super().__init__() in __init__.

2. Step 2: Verify forward method and layers

   class Discriminator(nn.Module): def __init__(self): super().__init__() self.fc = nn.Linear(784, 1) def forward(self, x): return torch.sigmoid(self.fc(x)) correctly defines the class with inheritance, super call, layer, and forward method.

3. Final Answer:

   Correct PyTorch discriminator class with inheritance and super call -> Option D

4. Quick Check:

   Proper nn.Module subclass = A [OK]

Hint: PyTorch models inherit nn.Module and call super().__init__() [OK]

Common Mistakes:

• Forgetting to inherit nn.Module
• Missing super().__init__() call
• Defining forward outside class

3. Given this PyTorch code snippet for a generator and discriminator, what will be the output shape of fake_output?

noise = torch.randn(16, 100)
generator = Generator()
fake_data = generator(noise)
discriminator = Discriminator()
fake_output = discriminator(fake_data)

Assume generator outputs shape (16, 784) and discriminator outputs shape (batch_size, 1).

medium

A. torch.Size([16, 1])

B. torch.Size([16, 784])

C. torch.Size([100, 16])

D. torch.Size([1, 16])

Solution

1. Step 1: Understand generator output shape

   The generator outputs fake data with shape (16, 784), meaning batch size 16 and 784 features.

2. Step 2: Understand discriminator input and output shape

   The discriminator takes input of shape (16, 784) and outputs shape (16, 1), one score per sample.

3. Final Answer:

   torch.Size([16, 1]) -> Option A

4. Quick Check:

   Discriminator output shape = (batch_size, 1) = A [OK]

Hint: Discriminator outputs one score per input sample [OK]

Common Mistakes:

• Confusing input and output shapes
• Mixing batch size and feature dimensions
• Assuming output shape matches noise size

4. Identify the error in this discriminator training step code snippet:

optimizer.zero_grad()
output = discriminator(fake_data)
loss = criterion(output, real_labels)
loss.backward()
optimizer.step()

Assuming fake_data is generated by the generator and real_labels are labels for real data.

medium

A. Using real_labels for fake_data loss calculation

B. Not calling optimizer.zero_grad() before backward

C. Using fake_data instead of real_data for discriminator input

D. Missing loss.backward() call

Solution

1. Step 1: Check input and labels for discriminator

   The discriminator should compare fake_data with fake labels, not real_labels.

2. Step 2: Identify label mismatch in loss calculation

   Using real_labels with fake_data causes wrong loss and training signals.

3. Final Answer:

   Using real_labels for fake_data loss calculation -> Option A

4. Quick Check:

   Fake data needs fake labels = C [OK]

Hint: Match fake data with fake labels during discriminator training [OK]

Common Mistakes:

• Mixing real and fake labels
• Skipping zero_grad() call
• Forgetting loss.backward()

5. You want to improve your GAN training stability. Which technique involving the discriminator can help reduce mode collapse and improve generator quality?

hard

A. Use dropout only in the generator

B. Use label smoothing on discriminator targets

C. Remove the discriminator and train generator alone

D. Increase generator learning rate only

Solution

1. Step 1: Understand mode collapse and training stability

   Mode collapse happens when generator produces limited outputs; discriminator tricks help prevent it.

2. Step 2: Identify label smoothing effect

   Label smoothing softens discriminator targets, preventing it from becoming too confident and stabilizing training.

3. Final Answer:

   Use label smoothing on discriminator targets -> Option B

4. Quick Check:

   Label smoothing stabilizes discriminator = D [OK]

Hint: Smooth discriminator labels to stabilize GAN training [OK]

Common Mistakes:

• Ignoring discriminator role in stability
• Changing only generator learning rate
• Removing discriminator breaks GAN