Train Your First GAN Model | Let‘s Talk About GANs Part 2

"Generative Adversarial Networks are the most interesting idea in machine learning in the last 10 years." — Yann LeCun, Chief AI Scientist at Facebook

In a previous article, we introduced the fundamental concepts behind Generative Adversarial Networks (GANs). To briefly recap, GANs are a class of deep learning models that can generate new data that mimics a given training dataset. They consist of two neural networks — a generator and a discriminator — that are trained simultaneously in a minimax game. The generator tries to fool the discriminator by generating realistic fake data, while the discriminator tries to distinguish between the real and fake samples. Through this adversarial process, the generator learns to create data that is virtually indistinguishable from the real thing.

Today, major tech companies are heavily investing in GAN research and development for a wide variety of applications:

• Adobe is using GANs to power next-generation Photoshop features
• Google is applying GANs to text generation
• IBM is leveraging GANs for data augmentation to train more robust classification models
• Snapchat and TikTok use GAN-based filters to transform users‘ pictures and videos
• Disney is using GANs to enhance video resolution for their movies

Clearly, GANs have immense potential to shape the future of creative content generation and machine learning as a whole. Are you ready to ride the wave of this groundbreaking technology? In this tutorial, we‘ll equip you with a practical understanding of how to code up and train your very own GAN model. Let‘s get started!

Discriminator Network

The job of the discriminator is to classify input data as either real (from the actual dataset) or fake (generated by the generator network). Essentially, it is just a standard convolutional neural network (CNN) trained as a binary classifier.

Here are the steps involved in training the discriminator:

1. Sample a mini-batch of m real examples {x(1), …, x(m)} from the dataset with corresponding labels y=1.

2. Sample a mini-batch of m fake examples {G(z(1)), …, G(z(m))} from the generator with corresponding labels y=0.

3. Train the discriminator network D on this combined mini-batch of 2m examples, adjusting its weights to minimize the binary cross-entropy loss:

where D(x) is the predicted probability that input x is real rather than fake.

By minimizing this loss, the discriminator learns to output values close to 1 for real examples and values close to 0 for fake examples.

In practice, the discriminator is usually implemented as a CNN with a sigmoid output layer to squash the raw logits into the range [0, 1]. Here‘s what the PyTorch code looks like:

class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 64, 4, 2, 1, bias=False)
        self.conv2 = nn.Conv2d(64, 128, 4, 2, 1, bias=False)
        self.bn2 = nn.BatchNorm2d(128)
        self.conv3 = nn.Conv2d(128, 256, 4, 2, 1, bias=False) 
        self.bn3 = nn.BatchNorm2d(256)
        self.conv4 = nn.Conv2d(256, 1, 4, 1, 0, bias=False)

    def forward(self, x):
        x = F.leaky_relu(self.conv1(x), 0.2)
        x = F.leaky_relu(self.bn2(self.conv2(x)), 0.2)
        x = F.leaky_relu(self.bn3(self.conv3(x)), 0.2)
        return torch.sigmoid(self.conv4(x))

This network takes in 28×28 grayscale MNIST images as input and outputs a single probability between 0 and 1. It has 4 convolutional layers with an increasing number of filters, batch normalization for stability, and leaky ReLU activations to allow gradients to flow backwards even when outputs are negative. The final layer collapses the output down to a single value which is passed through a sigmoid to obtain the predicted probability of the input being real.

Generator Network

The generator network takes random noise as input and upsamples it to generate fake data examples that ideally look indistinguishable from the real data. The goal is to fool the discriminator into classifying the generated samples as real.

The key steps for training the generator are:

1. Sample a mini-batch of m random noise vectors {z(1), …, z(m)} from a prior distribution, usually just a standard normal distribution.

2. Feed the noise through the generator network G to generate fake data samples {G(z(1)), …, G(z(m))}.

3. Feed the generated samples to the discriminator D to get predicted probabilities {D(G(z(1))), …, D(G(z(m)))}.

4. Train the generator network G to maximize the probability of the discriminator being mistaken, i.e. classifying fake samples as real. The loss function for the generator is:

Intuitively, the generator is trained to produce samples that the discriminator thinks are real, i.e. have a high D(G(z)). By maximizing this loss, the generator learns to generate increasingly realistic fake examples over time.

One common generator architecture is a transposed convolutional network that gradually upsamples low-resolution noise to a high-resolution image. Here‘s a simple 4-layer generator in PyTorch:

class Generator(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(100, 256*7*7)
        self.bn1 = nn.BatchNorm1d(256*7*7)
        self.deconv1 = nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False)  
        self.bn2 = nn.BatchNorm2d(128)
        self.deconv2 = nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False)
        self.bn3 = nn.BatchNorm2d(64)
        self.deconv3 = nn.ConvTranspose2d(64, 1, 4, 2, 1, bias=False)

    def forward(self, x):
        x = F.relu(self.bn1(self.fc1(x)))
        x = x.view(-1, 256, 7, 7)  # Resize to 256 x 7 x 7
        x = F.relu(self.bn2(self.deconv1(x)))
        x = F.relu(self.bn3(self.deconv2(x))) 
        x = torch.tanh(self.deconv3(x))
        return x

This generator takes a 100-dimensional noise vector as input, projects and reshapes it to a 7×7 spatial feature map, then applies a series of transposed convolutions with batch normalization and ReLU to gradually upsample to a final 28×28 grayscale image. The tanh output activation keeps pixel values normalized to the range [-1, 1].

Adversarial Training

To train the full GAN system, we alternate between taking gradient steps for the discriminator and the generator according to the following pseudo-code:

for number of training iterations do:

    # Train the discriminator for one step
    sample m examples {x(1), ..., x(m)} from dataset 
    sample m noise samples {z(1), ..., z(m)} from prior
    generate m fake samples {G(z(1)), ..., G(z(m))} from noise using G
    calculate discriminator loss on real and fake mini-batches
    update discriminator weights to minimize loss

    # Train the generator for one step
    sample m noise samples {z(1), ..., z(m)} from prior
    generate m fake samples {G(z(1)), ..., G(z(m))} from noise using G  
    calculate generator loss using discriminator predictions
    update generator weights to maximize loss

end

By training the networks in this alternating fashion, the discriminator learns to distinguish real and fake data while the generator learns to fool the discriminator. As training progresses, the generator produces increasingly realistic samples and the discriminator becomes increasingly good at flagging fakes. The end result is a generator network that can create novel data highly similar to the original training set.

Here‘s how we can implement GAN training in PyTorch using Adam optimizers:

# Initialize models and optimizers
generator = Generator()
discriminator = Discriminator()

g_optimizer = optim.Adam(generator.parameters(), lr=0.0002)
d_optimizer = optim.Adam(discriminator.parameters(), lr=0.0002)

# Loss function
criterion = nn.BCELoss()

# Training loop 
for epoch in range(num_epochs):
    for batch in dataloader:

        # Train discriminator
        discriminator.zero_grad()

        real_data = batch[0]
        batch_size = real_data.size(0)
        real_labels = torch.ones(batch_size, 1)

        real_preds = discriminator(real_data)
        d_real_loss = criterion(real_preds, real_labels)

        noise = torch.randn(batch_size, 100)
        fake_data = generator(noise)
        fake_labels = torch.zeros(batch_size, 1)

        fake_preds = discriminator(fake_data)
        d_fake_loss = criterion(fake_preds, fake_labels)

        d_loss = d_real_loss + d_fake_loss
        d_loss.backward()
        d_optimizer.step()

        # Train generator
        generator.zero_grad()

        noise = torch.randn(batch_size, 100)  
        fake_data = generator(noise)

        fake_preds = discriminator(fake_data)
        g_loss = criterion(fake_preds, real_labels)

        g_loss.backward()
        g_optimizer.step()

Some key things to note:

• We create separate optimizers for the generator and discriminator which independently update their respective weights based on their loss functions.
• When training the discriminator, we calculate losses for a mini-batch of real examples with target labels of 1 and a mini-batch of fake examples with labels of 0, then add the losses and backpropagate.
• When training the generator, we generate a mini-batch of fake examples and calculate the loss using the same target labels as the real examples. Backpropagating this loss encourages the generator to produce samples that fool the discriminator.
• We use label smoothing for the real examples when computing the discriminator loss. Instead of target values of exactly 1, we use a soft target like 0.9 which has been shown to improve training stability and reduce overconfidence.

Generating Samples

After training the GAN for some number of epochs, we can use the generator network to create new data samples from scratch. The process is simple — pass randomly sampled noise vectors through the generator and voila, brand new data that didn‘t exist before!

with torch.no_grad():
    noise = torch.randn(64, 100).to(device)  # 64 samples
    generated_images = generator(noise).cpu()

# Visualize generated images
plt.figure(figsize=(8, 8))
for i in range(64):
    plt.subplot(8, 8, i+1)
    plt.imshow(generated_images[i, 0], cmap=‘gray‘)
    plt.axis(‘off‘)
plt.tight_layout()
plt.show()  

Here we sample 64 noise vectors from a standard normal distribution, pass them through our trained generator, and plot the resulting digit images in an 8×8 grid.

The generated samples won‘t be perfect, but even rough approximations of the training data show the incredible power of GANs. With enough compute, data, and tuning, GANs can be used to generate photorealistic images, 3D models, music, speech, and more.

Hopefully this gives you a solid foundation to start building your own GAN models. The key things to remember are:

1. GANs consist of a generator that creates new data and a discriminator that classifies data as real or fake.
2. The networks are trained adversarially — the generator tries to fool the discriminator and the discriminator tries to accurately classify real and fake data.
3. Loss functions are based on binary cross-entropy. Discriminator loss depends on performance on both real and fake examples. Generator loss depends on discriminator predictions for fake examples.
4. Training alternates between discriminator and generator gradient steps to slowly improve both models simultaneously.

I encourage you to experiment with different datasets, model architectures, loss functions, and training techniques. GANs are still an active area of research with vast room for creativity and exploration. Go build something amazing!

References

1. Goodfellow et al. 2014. Generative Adversarial Networks. https://arxiv.org/abs/1406.2661
2. Arjovsky & Bottou. 2017. Towards Principled Methods for Training Generative Adversarial Networks. https://arxiv.org/abs/1701.04862
3. PyTorch DCGAN Tutorial. https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html