Simple Generative Adversarial Network (GAN) Pipeline

A step-by-step guide to building a basic Generative Adversarial Network (GAN) to generate new images.

Libraries Used

This pipeline uses a few key Python libraries. Here is a quick overview of their roles:

• tensorflow: The core deep learning framework for building the generator and discriminator.
• numpy: Used for handling numerical operations, particularly for creating random noise vectors.
• matplotlib: Used for visualizing the generated images.

1. Setup and Data Loading

• What the code does: It imports libraries and loads the MNIST dataset, which will be the source of "real" images for the discriminator.
• Behind the scenes: The GAN's job is to learn the distribution of the real images (in this case, handwritten digits) and generate new ones that look similar. The MNIST dataset is a simple starting point.
• Abstraction: The dataset loading process is handled by Keras, providing clean NumPy arrays for both training and evaluation.

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

# Load the MNIST dataset
(train_images, train_labels), (_, _) = tf.keras.datasets.mnist.load_data()

# Preprocess the data
train_images = train_images.reshape(train_images.shape[0], 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)

2. Building the Generator and Discriminator

• What the code does: It defines two separate models: a `generator` that takes random noise and produces an image, and a `discriminator` that takes an image and outputs a probability of it being "real."
• Behind the scenes: The generator uses `Conv2DTranspose` layers to "upsample" from a small latent space to a full-sized image. The discriminator uses standard `Conv2D` layers to downsample and classify the image.
• Abstraction: The `tf.keras.Sequential` API and the different layer types make it straightforward to define the architecture of these two competing neural networks.

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

    model.add(layers.Reshape((7, 7, 256)))
    
    model.add(layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh'))
    return model

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

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

    model.add(layers.Flatten())
    model.add(layers.Dense(1))
    return model

3. The Training Loop

• What the code does: This is the core of the GAN. It defines a custom training loop where the discriminator is first trained on both real and fake images, and then the generator is trained to fool the discriminator.
• Behind the scenes: This is a "game" where the two networks improve over time. The generator gets better at producing realistic images, and the discriminator gets better at spotting fakes. This adversarial process drives the training.
• Abstraction: This step requires a custom training loop (`@tf.function`) to manage the separate loss functions and gradient updates for the two models.

cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

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)

generator_optimizer = tf.keras.optimizers.Adam(1e-4)
discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)

# The main training step
@tf.function
def train_step(images):
    noise = tf.random.normal([BATCH_SIZE, 100])

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