Simple GAN implementation in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Simple GAN implementation

Problem:You want to train a simple Generative Adversarial Network (GAN) to generate images similar to handwritten digits from the MNIST dataset.

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

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

Your Task

Improve the GAN training so that the generator produces clearer, more realistic handwritten digit images after 50 epochs.

Use the MNIST dataset only.

Keep the basic GAN architecture but improve training stability.

Do not change the dataset or switch to a different GAN variant.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

TensorFlow

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

# Load and preprocess MNIST dataset
(train_images, _), (_, _) = tf.keras.datasets.mnist.load_data()
train_images = train_images.reshape(-1, 28, 28, 1).astype('float32')
train_images = (train_images - 127.5) / 127.5  # Normalize to [-1, 1]

BUFFER_SIZE = 60000
BATCH_SIZE = 256

train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)

# Generator model
def make_generator_model():
    model = tf.keras.Sequential([
        layers.Dense(7*7*256, use_bias=False, input_shape=(100,)),
        layers.BatchNormalization(),
        layers.LeakyReLU(),

        layers.Reshape((7, 7, 256)),

        layers.Conv2DTranspose(128, (5,5), strides=(1,1), padding='same', use_bias=False),
        layers.BatchNormalization(),
        layers.LeakyReLU(),

        layers.Conv2DTranspose(64, (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='tanh')
    ])
    return model

# Discriminator model with dropout and label smoothing

def make_discriminator_model():
    model = tf.keras.Sequential([
        layers.Conv2D(64, (5,5), strides=(2,2), padding='same', input_shape=[28,28,1]),
        layers.LeakyReLU(),
        layers.Dropout(0.3),

        layers.Conv2D(128, (5,5), strides=(2,2), padding='same'),
        layers.LeakyReLU(),
        layers.Dropout(0.3),

        layers.Flatten(),
        layers.Dense(1)
    ])
    return model

# Loss functions
cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

def discriminator_loss(real_output, fake_output):
    # Label smoothing for real labels
    real_labels = tf.ones_like(real_output) * 0.9
    real_loss = cross_entropy(real_labels, real_output)
    fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)
    total_loss = real_loss + fake_loss
    return total_loss

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

# Create models
generator = make_generator_model()
discriminator = make_discriminator_model()

# Optimizers with adjusted learning rates
generator_optimizer = tf.keras.optimizers.Adam(1e-4)
discriminator_optimizer = tf.keras.optimizers.Adam(5e-5)

# Training loop
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([BATCH_SIZE, 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

def train(dataset, epochs):
    for epoch in range(epochs):
        gen_loss_avg = 0
        disc_loss_avg = 0
        batch_count = 0
        for image_batch in dataset:
            gen_loss, disc_loss = train_step(image_batch)
            gen_loss_avg += gen_loss
            disc_loss_avg += disc_loss
            batch_count += 1
        gen_loss_avg /= batch_count
        disc_loss_avg /= batch_count
        print(f'Epoch {epoch+1}, Generator loss: {gen_loss_avg:.4f}, Discriminator loss: {disc_loss_avg:.4f}')

train(train_dataset, EPOCHS)

# Generate and show images after training
import matplotlib.pyplot as plt

noise = seed
generated_images = generator(noise, training=False)

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] * 127.5 + 127.5).numpy().astype('uint8'), cmap='gray')
    plt.axis('off')
plt.show()

Added dropout layers in the discriminator to reduce overfitting and overpowering.

Applied label smoothing to real labels in discriminator loss to stabilize training.

Reduced discriminator learning rate to slow down its learning relative to generator.

Kept generator learning rate higher to help it catch up.

Kept the basic GAN architecture but improved training stability with these changes.

Results Interpretation

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

After: Generator loss: ~0.45, Discriminator loss: ~0.65, images clearer and realistic.

Adding dropout and label smoothing helps balance GAN training by preventing the discriminator from overpowering the generator, leading to better quality generated images.

Bonus Experiment

Try training the GAN with a different optimizer like RMSprop or with different batch sizes to see how it affects training stability and image quality.

💡 Hint

Changing optimizer or batch size can affect how fast models learn and how stable training is. Experiment with smaller batch sizes or RMSprop optimizer with a low learning rate.

Practice

1. In a simple GAN, what is the main role of the generator?

easy

A. To classify data as real or fake

B. To create fake data that looks real

C. To optimize the loss function of the discriminator

D. To preprocess the input data

Solution

1. Step 1: Understand the generator's purpose

   The generator's job is to produce fake data samples that resemble real data.

2. Step 2: Differentiate from discriminator's role

   The discriminator's role is to distinguish real from fake data, not the generator's.

3. Final Answer:

   To create fake data that looks real -> Option B

4. Quick Check:

   Generator = create fake data [OK]

Hint: Generator makes fake data to fool discriminator [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 the input noise vector for the generator in TensorFlow?

easy

A. noise = tf.random.normal([batch_size, noise_dim])

B. noise = tf.zeros([batch_size, noise_dim])

C. noise = tf.ones([batch_size, noise_dim])

D. noise = tf.random.uniform([noise_dim])

Solution

1. Step 1: Identify typical noise input shape

   The noise vector usually has shape [batch_size, noise_dim] to generate a batch of samples.

2. Step 2: Check correct TensorFlow function

   tf.random.normal generates random normal noise, commonly used for GAN input.

3. Final Answer:

   noise = tf.random.normal([batch_size, noise_dim]) -> Option A

4. Quick Check:

   Use tf.random.normal with batch size and noise dim [OK]

Hint: Use tf.random.normal with batch and noise dims [OK]

Common Mistakes:

• Using tf.zeros or tf.ones which produce constant noise
• Missing batch dimension in noise shape
• Using tf.random.uniform without batch dimension

3. Given this TensorFlow code snippet for a GAN training step, what will be the shape of fake_output?

noise = tf.random.normal([32, 100])
fake_images = generator(noise)
fake_output = discriminator(fake_images)

medium

A. [32, 1]

B. [100, 32]

C. [32, 100]

D. [1, 32]

Solution

1. Step 1: Analyze input batch size

   The noise input has batch size 32, so generator outputs 32 fake images.

2. Step 2: Understand discriminator output shape

   The discriminator outputs a prediction per image, usually shape [batch_size, 1].

3. Final Answer:

   [32, 1] -> Option A

4. Quick Check:

   Discriminator output shape = [batch_size, 1] [OK]

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

Common Mistakes:

• Confusing noise dimension with output shape
• Swapping batch size and feature dimensions
• Assuming discriminator outputs full image shape

4. You have this error during GAN training: ValueError: Shapes (32, 1) and (64, 1) are incompatible. What is the most likely cause?

medium

A. Discriminator model is missing an output layer

B. Generator output shape is incorrect

C. Mismatch between batch sizes of real and fake data during training

D. Noise vector dimension is too large

Solution

1. Step 1: Understand the error message

   The error shows two tensors with batch sizes 32 and 64 being compared or combined.

2. Step 2: Identify cause in GAN training

   Real and fake batches must have the same batch size when computing loss or concatenating.

3. Final Answer:

   Mismatch between batch sizes of real and fake data during training -> Option C

4. Quick Check:

   Batch size mismatch causes shape errors [OK]

Hint: Ensure real and fake batches have same size [OK]

Common Mistakes:

• Ignoring batch size mismatch
• Blaming generator output shape without checking batch
• Assuming model layers missing cause shape error

5. You want to improve your simple GAN's training stability. Which change is best to try first?

hard

A. Train discriminator twice as often as generator

B. Increase noise vector dimension to 1000

C. Remove batch normalization layers from generator

D. Use label smoothing on discriminator labels

Solution

1. Step 1: Understand training stability issues

   GANs often suffer from unstable training causing oscillations or mode collapse.

2. Step 2: Evaluate common techniques

   Label smoothing (e.g., using 0.9 instead of 1 for real labels) helps discriminator not be too confident, improving stability.

3. Final Answer:

   Use label smoothing on discriminator labels -> Option D

4. Quick Check:

   Label smoothing improves GAN training stability [OK]

Hint: Try label smoothing first to stabilize GAN training [OK]

Common Mistakes:

• Increasing noise dimension too much adds complexity
• Removing batch norm can hurt generator learning
• Training discriminator more often can cause imbalance