Implementing Conditional GANs in PyTorch for Controlled Synthesis

Generative Adversarial Networks (GANs) have gained popularity for their ability to generate realistic datasets by training two models simultaneously: a generator and a discriminator. Conditional GANs (cGANs) are an extension of this concept which allow us to control the data generated by conditioning both the generator and discriminator on some extra information, such as class labels.

In this article, we will walk through the steps required to implement Conditional GANs using PyTorch, a popular deep learning library. Let’s start by setting up our environment.

Environment Setup

First, ensure that you have Python and PyTorch installed. You can install PyTorch by following the instructions on their official website. Additionally, we will be using torchvision for datasets and utilities, which can be installed via pip:

pip install torchvision

Dataset Preparation

For our cGAN implementation, we will use the MNIST dataset which consists of labeled digit images.

import torch
import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

trainset = torchvision.datasets.MNIST(
    root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(
    trainset, batch_size=64, shuffle=True)

Building the Generator and Discriminator

The generator and discriminator networks in a cGAN are similar to those in a regular GAN, but they are conditioned on label data. We will incorporate this conditional data by concatenating it with any input received by the networks.

import torch.nn as nn

class Generator(nn.Module):
    def __init__(self, nz, num_classes, ngf, image_size):
        super(Generator, self).__init__()
        self.embed = nn.Embedding(num_classes, num_classes)
        self.model = nn.Sequential(
            nn.Linear(nz + num_classes, ngf * 4),
            nn.ReLU(True),
            nn.Linear(ngf * 4, ngf * 8),
            nn.ReLU(True),
            nn.Linear(ngf * 8, image_size * image_size),
            nn.Tanh()  
        )

    def forward(self, noise, labels):
        label_embedding = self.embed(labels)
        gen_input = torch.cat((noise, label_embedding), -1)
        return self.model(gen_input).view(-1, 1, 28, 28)

class Discriminator(nn.Module):
    def __init__(self, ndf, num_classes, image_size):
        super(Discriminator, self).__init__()
        self.embed = nn.Embedding(num_classes, image_size * image_size)
        self.model = nn.Sequential(
            nn.Linear(image_size * image_size + image_size * image_size, ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(ndf * 4, ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(ndf * 2, 1),
            nn.Sigmoid()
        )

    def forward(self, img, labels):
        label_embedding = self.embed(labels).view(labels.size(0), -1)
        disc_input = torch.cat((img.view(img.size(0), -1), label_embedding), -1)
        return self.model(disc_input)

Training the Conditional GAN

With the generator and discriminator ready, we can proceed to train the network. Key to GAN training is the adversarial process where both models improve themselves towards fooling each other.

import torch.optim as optim

# Hyperparameters
num_epochs = 100
nz = 100  # size of the latent z vector
num_classes = 10
lr = 0.0002
beta1 = 0.5

# Instantiating models
netG = Generator(nz, num_classes, ngf=64, image_size=28).to(device)
netD = Discriminator(ndf=64, num_classes=num_classes, image_size=28).to(device)

# Optimizers
optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))

# Training Loop
for epoch in range(num_epochs):
    for i, data in enumerate(trainloader, 0):
        # Training code for Discriminator and Generator
        # as discussed in standard GAN workflows
        pass  # Provided just for code structural context

1. Generate noise vectors and concatenate them with labels for the generator.
2. Generate fake images using the conditioned generator.
3. Compute loss of discriminator on both real and fake data independently.
4. Backpropagate the discriminator loss and update its weights accordingly.
5. Calculate generator loss and backpropagate errors through it.
6. Finally, update the generator's weights.

Enjoy experimenting further with this implementation by trying out different initial conditions such as varying shapes, network sizes, or training times. cGANs provide tremendous flexibility and control over generated outputs which is powerful across numerous applications.