Generating Realistic Images with GANs in TensorFlow

Introduction

In recent years, generative adversarial networks (GANs) have emerged as one of the most exciting developments in deep learning and computer vision. GANs are a class of machine learning models that can generate new, realistic data—such as images—that mimic the properties of a training dataset. The potential applications are vast, from generating photorealistic images and artwork to improving data efficiency for other deep learning tasks.

In this post, we‘ll dive into how GANs work and walk through an example of building one in TensorFlow to generate images of handwritten digits. By the end, you‘ll understand the key components of a GAN and how to implement one yourself. Let‘s get started!

What are Generative Adversarial Networks?

At their core, GANs are made up of two competing neural networks—a generator and a discriminator—that engage in a min-max game during training:

The generator takes random noise as input and attempts to generate realistic "fake" data (e.g. images) that can fool the discriminator. Its goal is to maximize the probability that the discriminator classifies its outputs as real.

The discriminator takes real data and generated fake data as input and attempts to distinguish between the two. Its goal is to not be fooled by the increasingly realistic fakes coming from the generator.

Over the course of many training iterations, the generator learns to map from the latent noise space to the distribution of real data, getting better and better at generating realistic samples. Meanwhile, the discriminator gets better at correctly classifying real vs fake. In this way, the two networks push each other to improve, converging on an equilibrium where the generator produces samples that the discriminator can no longer distinguish from real data.

The key innovation of GANs is this adversarial training paradigm. Unlike other generative models that explicitly maximize the likelihood of the training data, GANs learn the distribution implicitly through the discriminator. This allows them to generate sharp, realistic samples across diverse domains.

GANs in TensorFlow

TensorFlow is one of the most popular deep learning frameworks, and it provides all the building blocks needed to implement state-of-the-art GANs. Let‘s walk through an example of building a simple GAN to generate MNIST digits.

We‘ll start by importing the required libraries and loading the MNIST dataset:

import tensorflow as tf
from tensorflow.keras.datasets import mnist

(train_images, train_labels), (_, _) = 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 to [-1, 1]

Next, we‘ll define the architectures for the generator and discriminator. The generator will take a 100-dimensional random noise vector as input and output a 28×28 grayscale image. It does this by passing the noise through a series of upsampling conv layers:

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

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

        tf.keras.layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding=‘same‘, use_bias=False),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.LeakyReLU(),

        tf.keras.layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding=‘same‘, use_bias=False, activation=‘tanh‘)
    ])

    return model

The discriminator takes 28×28 images as input (either real or generated) and outputs a single probability value between 0-1, classifying the input as real or fake. It does this with a series of conv layers that downsample the input:

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

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

        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(1)
    ])

    return model  

With the generator and discriminator defined, we can instantiate them and set up the loss functions and optimizers. We‘ll use binary cross entropy loss and the Adam optimizer:

generator = make_generator_model()
discriminator = make_discriminator_model()

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)

Finally, we can define the training loop. For each batch of real images, we‘ll:

1. Generate fake images from random noise
2. Train the discriminator on the real and fake images
3. Train the generator by trying to fool the discriminator with its fake images

We‘ll also generate and save sample images every 100 batches to visualize the generator‘s progress:

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

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

        generate_and_save_images(generator,
                                 epoch + 1,
                                 seed)

def generate_and_save_images(model, epoch, test_input):
    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()

train(train_dataset, EPOCHS)    

Let‘s unpack what‘s happening here. In each train_step, we:

1. Sample random noise and generate fake images from it using the generator
2. Get the discriminator‘s predictions on the real and fake images
3. Calculate the loss for the generator and discriminator
4. Get the gradients of the loss w.r.t. the model parameters
5. Update the model parameters using the optimizer

By repeating this process over many epochs, the discriminator learns to classify real vs fake images, while the generator learns to fool the discriminator. After 50 epochs of training on MNIST, the generator has learned to produce quite plausible looking handwritten digits:

We can see that the digits start out as unrecognizable noise and become more and more realistic over time as the generator improves, until they are largely indistinguishable from real MNIST images to the human eye. This showcases the power of the adversarial training paradigm – the generator has implicitly learned the data distribution well enough to generate novel samples that match the properties of real data from the domain.

Taking it Further

This is just a simple example to illustrate the core concepts of GANs. In practice, modern GAN architectures for generating high resolution images are far more complex, involving techniques like:

• Progressively growing the model to learn features at multiple scales
• Using conditional input to control attributes of generated images
• Modified loss functions to stabilize training and avoid mode collapse
• Regularization techniques like spectral normalization

Popular architectures like StyleGAN can generate photorealistic images of faces, objects, and scenes at resolutions of 1024×1024 pixels and beyond. The quality is staggering – in many cases generated images are very difficult to distinguish from real photos.

The potential applications are endless. GANs can be used for creative purposes like generating art, super-resolution (increasing the resolution of images), or editing the style and content of images and videos. They can also augment datasets for downstream supervised learning tasks.

However, like any powerful technology, GANs also raise important ethical considerations around potential misuse, such as generating fake media for disinformation. As the technology advances, it will be crucial to develop it responsibly.

Conclusion

GANs are a powerful class of generative models that can produce amazingly realistic data across domains. Their adversarial training paradigm allows them to implicitly learn rich data distributions and generate novel samples that match the properties of real data.

Hopefully this post gave you a taste of how GANs work and how to get started building them in TensorFlow. The field is progressing incredibly quickly, with bigger and better models coming out all the time that push the boundaries of what‘s possible.

While there are valid concerns around potential misuse, GANs also have a huge potential to be used for good. From creating art to accelerating and democratizing ML research, it‘s an exciting technology to keep an eye on. Thanks for reading!