GAN for image generation in Computer Vision - ML Experiment: Train & Evaluate

Experiment - GAN for image generation

Problem:Generate realistic handwritten digit images using a GAN trained on the MNIST dataset.

Current Metrics:Training loss generator: 1.2, Training loss discriminator: 0.3, Generated images are blurry and lack detail.

Issue:The GAN model is not producing sharp images; the generator and discriminator losses are imbalanced, indicating unstable training and mode collapse.

Your Task

Improve the GAN training to generate clearer, more realistic handwritten digit images with balanced generator and discriminator losses.

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

Use the MNIST dataset only.

Do not increase training epochs beyond 50.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

Computer Vision

import tensorflow as tf
from tensorflow.keras import layers
import numpy as np
import matplotlib.pyplot as plt

# Load and preprocess MNIST data
def load_mnist():
    (x_train, _), (_, _) = tf.keras.datasets.mnist.load_data()
    x_train = x_train.astype('float32') / 255.0
    x_train = np.expand_dims(x_train, axis=-1)  # Shape (60000, 28, 28, 1)
    return x_train

# Build generator model
def build_generator():
    model = tf.keras.Sequential([
        layers.Dense(7*7*128, use_bias=False, input_shape=(100,)),
        layers.BatchNormalization(),
        layers.LeakyReLU(),
        layers.Reshape((7, 7, 128)),
        layers.Conv2DTranspose(64, (5,5), strides=(1,1), padding='same', use_bias=False),
        layers.BatchNormalization(),
        layers.LeakyReLU(),
        layers.Conv2DTranspose(32, (5,5), strides=(2,2), padding='same', use_bias=False),
        layers.BatchNormalization(),
        layers.LeakyReLU(),
        layers.Conv2DTranspose(1, (5,5), strides=(2,2), padding='same', use_bias=False, activation='sigmoid')
    ])
    return model

# Build discriminator model
def build_discriminator():
    model = tf.keras.Sequential([
        layers.Conv2D(64, (5,5), strides=(2,2), padding='same', input_shape=[28,28,1]),
        layers.LeakyReLU(alpha=0.2),
        layers.Dropout(0.3),
        layers.Conv2D(128, (5,5), strides=(2,2), padding='same'),
        layers.LeakyReLU(alpha=0.2),
        layers.Dropout(0.3),
        layers.Flatten(),
        layers.Dense(1, activation='sigmoid')
    ])
    return model

# Loss functions and optimizers
cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=False)

def discriminator_loss(real_output, fake_output):
    real_loss = cross_entropy(tf.ones_like(real_output), real_output)
    fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)
    return real_loss + fake_loss

def generator_loss(fake_output):
    return cross_entropy(tf.ones_like(fake_output), fake_output)

# Prepare dataset
x_train = load_mnist()
batch_size = 256
train_dataset = tf.data.Dataset.from_tensor_slices(x_train).shuffle(60000).batch(batch_size)

# Create models
generator = build_generator()
discriminator = build_discriminator()

# Optimizers with different learning rates
generator_optimizer = tf.keras.optimizers.Adam(learning_rate=0.0002, beta_1=0.5)
discriminator_optimizer = tf.keras.optimizers.Adam(learning_rate=0.0001, beta_1=0.5)

# Training loop
import time

epochs = 50
noise_dim = 100
num_examples_to_generate = 16

seed = tf.random.normal([num_examples_to_generate, noise_dim])

@tf.function
def train_step(images):
    noise = tf.random.normal([tf.shape(images)[0], noise_dim])

    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
        generated_images = generator(noise, training=True)

        real_output = discriminator(images, training=True)
        fake_output = discriminator(generated_images, training=True)

        gen_loss = generator_loss(fake_output)
        disc_loss = discriminator_loss(real_output, fake_output)

    gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
    gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)

    generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))

    return gen_loss, disc_loss

# Training function

def train(dataset, epochs):
    for epoch in range(epochs):
        start = time.time()

        gen_loss_avg = 0
        disc_loss_avg = 0
        batches = 0

        for image_batch in dataset:
            gen_loss, disc_loss = train_step(image_batch)
            gen_loss_avg += gen_loss
            disc_loss_avg += disc_loss
            batches += 1

        gen_loss_avg /= batches
        disc_loss_avg /= batches

        print(f'Epoch {epoch+1}, Gen Loss: {gen_loss_avg:.4f}, Disc Loss: {disc_loss_avg:.4f}, Time: {time.time()-start:.2f} sec')

train(train_dataset, epochs)

# Generate and save images
import os
os.makedirs('generated_images', exist_ok=True)

noise = tf.random.normal([num_examples_to_generate, noise_dim])
generated_images = generator(noise, training=False)

import matplotlib.pyplot as plt

fig = plt.figure(figsize=(4,4))
for i in range(generated_images.shape[0]):
    plt.subplot(4, 4, i+1)
    plt.imshow(generated_images[i, :, :, 0], cmap='gray')
    plt.axis('off')
plt.savefig('generated_images/sample.png')
plt.close()

Added batch normalization layers in the generator to stabilize training.

Used LeakyReLU activation in both generator and discriminator for better gradient flow.

Added dropout layers in the discriminator to reduce overfitting.

Set different learning rates for generator (0.0002) and discriminator (0.0001) optimizers.

Kept training epochs to 50 to avoid overtraining.

Fixed batch size usage in train_step to handle last batch size correctly.

Results Interpretation

Before: Generator loss: 1.2, Discriminator loss: 0.3, Images blurry and unstable training.

After: Generator loss: ~0.45, Discriminator loss: ~0.55, Images sharper and realistic, stable training.

Adding batch normalization, dropout, and adjusting learning rates helps stabilize GAN training and reduces overfitting, resulting in better image quality.

Bonus Experiment

Try using a Wasserstein GAN (WGAN) with gradient penalty to further improve image quality and training stability.

💡 Hint

Replace the loss functions with Wasserstein loss and add gradient penalty to the discriminator loss.

Practice

1. What are the two main parts of a GAN used for image generation?

easy

A. Convolution and Pooling

B. Encoder and Decoder

C. Input Layer and Output Layer

D. Generator and Discriminator

Solution

1. Step 1: Understand GAN components

   A GAN consists of two parts: a generator that creates images and a discriminator that evaluates them.

2. Step 2: Match parts to options

   Only Generator and Discriminator correctly names these two parts as Generator and Discriminator.

3. Final Answer:

   Generator and Discriminator -> Option D

4. Quick Check:

   GAN parts = Generator and Discriminator [OK]

Hint: Remember GAN = Generator + Discriminator [OK]

Common Mistakes:

• Confusing GAN parts with autoencoder parts
• Thinking GAN uses Encoder and Decoder
• Mixing up layers with GAN components

2. Which of the following is the correct way to define the generator model in PyTorch for a GAN?

easy

A. generator = nn.Sequential(nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 100), nn.Softmax())

B. generator = nn.Sequential(nn.Conv2d(1, 64, 3), nn.Sigmoid())

C. generator = nn.Sequential(nn.Linear(100, 256), nn.ReLU(), nn.Linear(256, 784), nn.Tanh())

D. generator = nn.Sequential(nn.ConvTranspose2d(100, 1, 3), nn.ReLU())

Solution

1. Step 1: Identify generator input and output sizes

   The generator takes a noise vector (size 100) and outputs an image vector (size 784 for 28x28 images).

2. Step 2: Check model layers and activations

   generator = nn.Sequential(nn.Linear(100, 256), nn.ReLU(), nn.Linear(256, 784), nn.Tanh()) correctly uses Linear layers from 100 to 256 to 784 with ReLU and Tanh activations, common in GAN generators.

3. Final Answer:

   generator = nn.Sequential(nn.Linear(100, 256), nn.ReLU(), nn.Linear(256, 784), nn.Tanh()) -> Option C

4. Quick Check:

   Generator input=100, output=784 with ReLU and Tanh [OK]

Hint: Generator maps noise to image size with Linear + Tanh [OK]

Common Mistakes:

• Using Conv2d instead of Linear for generator input
• Reversing input-output sizes
• Using Softmax instead of Tanh at output

3. Given this PyTorch snippet for a GAN training step, what will be the shape of the generated images tensor?

noise = torch.randn(16, 100)
generated_images = generator(noise)

medium

A. (16, 784)

B. (16, 100)

C. (100, 16)

D. (784, 16)

Solution

1. Step 1: Understand input noise shape

   The noise tensor has shape (16, 100), meaning batch size 16 and noise vector length 100.

2. Step 2: Check generator output shape

   The generator maps each noise vector of size 100 to an image vector of size 784, so output shape is (16, 784).

3. Final Answer:

   (16, 784) -> Option A

4. Quick Check:

   Batch size 16, output image size 784 [OK]

Hint: Output shape matches batch size and image vector size [OK]

Common Mistakes:

• Confusing input noise shape with output shape
• Swapping batch size and feature dimensions
• Assuming output shape is (16, 100)

4. You have this error during GAN training: RuntimeError: size mismatch in Linear layer. What is the most likely cause?

medium

A. The input noise vector size does not match the first Linear layer input size

B. The discriminator output is not a probability

C. The generator uses ReLU instead of Tanh

D. The batch size is too large

Solution

1. Step 1: Understand size mismatch error

   This error usually means the input tensor size does not match the expected input size of a Linear layer.

2. Step 2: Identify likely cause in GAN

   If the noise vector size does not match the first Linear layer's input features in the generator, this error occurs.

3. Final Answer:

   The input noise vector size does not match the first Linear layer input size -> Option A

4. Quick Check:

   Size mismatch = input noise size error [OK]

Hint: Check noise vector size matches Linear layer input [OK]

Common Mistakes:

• Blaming activation functions for size errors
• Thinking batch size causes size mismatch
• Ignoring input dimension mismatch

5. You want to improve the quality of images generated by your GAN. Which of these strategies is most effective?

hard

A. Increase the size of the noise vector and train longer

B. Use a deeper generator and discriminator with batch normalization

C. Remove the discriminator and train only the generator

D. Use random labels instead of real/fake labels for training

Solution

1. Step 1: Consider model architecture improvements

   Deeper networks with batch normalization help stabilize training and improve image quality.

2. Step 2: Evaluate other options

   Increasing noise size alone or removing discriminator harms quality; random labels confuse training.

3. Final Answer:

   Use a deeper generator and discriminator with batch normalization -> Option B

4. Quick Check:

   Better architecture + batch norm = improved GAN images [OK]

Hint: Better models + batch norm improve GAN image quality [OK]

Common Mistakes:

• Thinking bigger noise vector alone improves quality
• Removing discriminator to simplify training
• Using random labels instead of correct ones