close

Articles Home » Programming Languages » Implementing GANs (Generative Adversarial Networks) from Scratch

Implementing GANs (Generative Adversarial Networks) from Scratch

Generative Adversarial Networks (GANs) have revolutionized the field of artificial intelligence by enabling machines to create data that is nearly indistinguishable from real data. From generating realistic images to creating synthetic data for training, GANs have a wide range of applications. This tutorial will guide you through the process of implementing GANs from scratch using Python and TensorFlow.

Prerequisites

Before diving into the implementation, it’s essential to have a good grasp of the following concepts:

  1. Neural Networks and Deep Learning basics.
  2. Python programming.
  3. TensorFlow or PyTorch basics.

This tutorial assumes that you are comfortable with these prerequisites.

Introduction to GANs

GANs, introduced by Ian Goodfellow and his colleagues in 2014, consist of two neural networks: the Generator and the Discriminator. These two networks compete against each other in a game-theoretic framework. The Generator tries to create realistic data, while the Discriminator attempts to distinguish between real and fake data. Over time, both networks improve, resulting in a Generator capable of producing highly realistic data.

Basic Structure of GANs

  1. Generator: Takes random noise as input and generates data.
  2. Discriminator: Takes data (either real or generated) as input and outputs a probability that the data is real.

The Generator and Discriminator are trained together in a two-player minimax game:

  • The Discriminator is trained to maximize the probability of assigning the correct label to both real and generated data.
  • The Generator is trained to minimize the probability that the Discriminator assigns the correct label to the generated data.

Mathematically, the objective can be expressed as:

$latex \min_G \max_D V(D, G) = \mathbb{E}{x \sim p{\text{data}}(x)} [\log D(x)] + \\
\mathbb{E}_{z \sim p_z(z)} [\log(1 – D(G(z)))]&s=2$

Where:

  • G is the Generator.
  • D is the Discriminator.
  • p_{\text{data}}(x) is the distribution of real data.
  • p_z(z) is the distribution of the noise input to the Generator.

Step-by-Step Implementation

We will implement a simple GAN to generate handwritten digits similar to those in the MNIST dataset. The MNIST dataset contains 28×28 grayscale images of handwritten digits (0-9).

Step 1: Setup and Dependencies

First, we need to install and import the necessary libraries. We will use TensorFlow for building and training our GAN.

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

# Ensure we are using TensorFlow 2.x
assert tf.__version__.startswith('2')

# Set random seed for reproducibility
np.random.seed(42)
tf.random.set_seed(42)Code language: Python (python)

Step 2: Data Preparation

We will use the MNIST dataset, which is available in TensorFlow Datasets. The dataset will be normalized to have values between -1 and 1, which is a common practice when training GANs.

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

# Batch and shuffle the data
BUFFER_SIZE = 60000
BATCH_SIZE = 256

train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)Code language: Python (python)

Step 3: Build the Generator

The Generator takes random noise as input and produces an image. We’ll use a simple neural network with Dense and Conv2DTranspose layers to generate 28×28 images.

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)))
    assert model.output_shape == (None, 7, 7, 256)  # Note: None is the batch size

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

    model.add(layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False))
    assert model.output_shape == (None, 14, 14, 64)
    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'))
    assert model.output_shape == (None, 28, 28, 1)

    return model

generator = make_generator_model()Code language: Python (python)

Step 4: Build the Discriminator

The Discriminator takes an image as input and outputs a scalar indicating whether the image is real or fake. We’ll use a convolutional neural network for this task.

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

discriminator = make_discriminator_model()Code language: Python (python)

Step 5: Define the Loss and Optimizers

The loss functions for the Generator and Discriminator are crucial for the training process. We will use binary cross-entropy loss for both. The Generator tries to minimize -log(D(G(z))), while the Discriminator tries to maximize log(D(x)) + log(1 - D(G(z))).

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)
    total_loss = real_loss + fake_loss
    return total_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)Code language: Python (python)

Step 6: Define the Training Loop

The training loop involves generating random noise, using it to produce fake images with the Generator, and then training both the Generator and Discriminator. We will also periodically generate and save images to monitor the progress of the training.

EPOCHS = 50
noise_dim = 100
num_examples_to_generate = 16

# We will reuse this seed over time (so it's easier to visualize progress in the GIF)
seed = tf.random.normal([num_examples_to_generate, noise_dim])

# Notice the use of `tf.function`
# This annotation causes the function to be "compiled".
@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))

def train(dataset, epochs):
    for epoch in range(epochs):
        for image_batch in dataset:
            train_step(image_batch)

        # Produce images for the GIF as we go
        display.clear_output(wait=True)
        generate_and_save_images(generator, epoch + 1, seed)

        # Save the model every 15 epochs
        if (epoch + 1) % 15 == 0:
            checkpoint.save(file_prefix = checkpoint_prefix)

    # Generate after the final epoch
    display.clear_output(wait=True)
    generate_and_save_images(generator, epochs, seed)

def generate_and_save_images(model, epoch, test_input):
    # Notice `training` is set to False.
    # This is so all layers run in inference mode (batchnorm).
    predictions = model(test_input, training=False)

    fig = plt.figure(figsize=(4, 4))

    for i in range(predictions.shape[0]):
        plt.subplot

(4, 4, i+1)
        plt.imshow(predictions[i, :, :, 0] * 127.5 + 127.5, cmap='gray')
        plt.axis('off')

    plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))
    plt.show()Code language: Python (python)

Step 7: Train the GAN

Now, we can train our GAN on the MNIST dataset.

train(train_dataset, EPOCHS)Code language: Python (python)

Visualizing the Results

Once the training is complete, we can visualize the generated images and see how the Generator has improved over time.

def display_image(epoch_no):
    return PIL.Image.open('image_at_epoch_{:04d}.png'.format(epoch_no))

display_image(EPOCHS)Code language: Python (python)

This concludes our implementation of a simple GAN from scratch. We have walked through the process of building, training, and visualizing a GAN using TensorFlow. While this is a basic implementation, it provides a solid foundation for exploring more advanced GAN architectures and techniques.

Advanced Topics

Conditional GANs (cGANs)

Conditional GANs extend the standard GAN framework by conditioning both the Generator and Discriminator on some extra information. This information can be anything from class labels to other modalities. For example, in a cGAN, the Generator could be conditioned on a class label to generate images of a specific class.

DCGANs (Deep Convolutional GANs)

Deep Convolutional GANs (DCGANs) are a popular and effective GAN architecture that replaces fully connected layers with convolutional layers. DCGANs have been shown to produce high-quality images and are widely used in practice. The architecture and training techniques introduced in DCGANs can be applied to various GAN variants.

Wasserstein GANs (WGANs)

Wasserstein GANs (WGANs) improve the training stability of GANs by using the Earth Mover’s (Wasserstein) distance instead of the Jensen-Shannon divergence. This modification leads to a more stable training process and better quality of generated samples. WGANs require a slight modification to the Discriminator (called the Critic in WGANs) and the use of weight clipping.

GAN Applications

GANs have numerous applications, including:

  1. Image Generation: Creating realistic images from random noise.
  2. Image-to-Image Translation: Translating images from one domain to another, such as turning sketches into photographs.
  3. Data Augmentation: Generating synthetic data to augment training datasets.
  4. Super-Resolution: Enhancing the resolution of images.
  5. Text-to-Image Synthesis: Generating images from textual descriptions.

Tips for Training GANs

  1. Use Batch Normalization: Helps stabilize training.
  2. Use LeakyReLU Activation: Helps gradients flow through the network.
  3. Label Smoothing: Smooth the labels for real and fake images to help the Discriminator generalize better.
  4. Avoid Mode Collapse: Monitor and adjust training to avoid the Generator producing limited varieties of outputs.
  5. Use Progressive Training: Start with a low resolution and gradually increase the resolution of generated images.

Conclusion

Implementing GANs from scratch provides a deep understanding of their underlying mechanics and challenges. While we have covered the basics and some advanced topics, the field of GANs is rapidly evolving with continuous research and new techniques. Experimenting with different architectures, loss functions, and training strategies will help you gain more insights and potentially contribute to the development of more robust and versatile GAN models.

Related posts: