Training a Wasserstein GAN (WGAN) in PyTorch for Stable Generative Results

In the realm of generative models, Generative Adversarial Networks (GANs) have established themselves as groundbreaking due to their potential in generating high-quality data. However, conventional GANs often face instability during training, leading to mode collapse or gradient vanishment issues. To tackle these problems, the Wasserstein GAN (WGAN) was introduced, which provides a more reliable cost function through the use of the Wasserstein (Earth Mover's) distance. In this article, we'll walk through how to implement and train a WGAN using PyTorch.

Background

The key innovation in WGAN is its revised loss function which relies on the earth mover's (Wasserstein) distance as a metric. This introduces smoother gradients and a better signal for the generator to learn from. The following are the core principles of WGAN:

• Use of the Wasserstein distance instead of Jensen-Shannon divergence.
• Enforcing a Lipschitz constraint on the critic by clipping its weights to a fixed capping value.
• Training the critic more than the generator to provide better gradient direction.

Setting Up the Environment

We’ll start by setting up our development environment. Ensure you have PyTorch installed alongside some usual makeshift data-handling libraries like NumPy and Matplotlib.

pip install torch torchvision numpy matplotlib

Building the WGAN Components

Now, let's build the fundamental components: the generator, the critic (or discriminator in classical GANs), and the training loop. We'll make use of PyTorch's Module class to structure our network.

Define the Generator

import torch
import torch.nn as nn

class Generator(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(Generator, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, output_dim),
            nn.Tanh()
        )
    
    def forward(self, x):
        return self.net(x)

Define the Critic

class Critic(nn.Module):
    def __init__(self, input_dim):
        super(Critic, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 128),
            nn.LeakyReLU(0.2),
            nn.Linear(128, 1)
        )

    def forward(self, x):
        return self.net(x)

Training the WGAN

The critical difference when training a WGAN is its weight clipping mechanism, as well as differing update rules compared to traditional GANs.

import itertools

# Hyperparameters
latent_dim = 100
output_dim = 784  # Example for MNIST
epochs = 100000
n_critic = 5
clip_value = 0.01

# Initialize generator and critic
generator = Generator(latent_dim, output_dim)
critic = Critic(output_dim)

# Optimizers
opt_gen = torch.optim.RMSprop(generator.parameters(), lr=0.00005)
opt_critic = torch.optim.RMSprop(critic.parameters(), lr=0.00005)

for epoch in range(epochs):
    for _ in range(n_critic):
        # Train critic
        opt_critic.zero_grad()
        # Sample real and fake data
        z = torch.randn(batch_size, latent_dim)
        real_data = get_real_data(batch_size)  # Implement this to bring real data samples
        fake_data = generator(z)
        
        # Calculate critic loss
        loss_critic = torch.mean(critic(fake_data)) - torch.mean(critic(real_data))
        loss_critic.backward()
        opt_critic.step()
        
        # Clip weights
        for p in critic.parameters():
            p.data.clamp_(-clip_value, clip_value)

    # Train generator
    opt_gen.zero_grad()
    z = torch.randn(batch_size, latent_dim)
    fake_data = generator(z)
    loss_gen = -torch.mean(critic(fake_data))
    loss_gen.backward()
    opt_gen.step()

    # Output running info
    if epoch % 100 == 0:
        print(f'Epoch [{epoch}/{epochs}] Loss D: {loss_critic.item()}, Loss G: {loss_gen.item()}')

As illustrated, the critic is trained multiple times (as defined by n_critic) for every generator update, promoting a robust learning criterion for the generator. Additionally, the weights of the critic are clipped to ensure the Lipschitz constraint.

Conclusion

By applying these modifications and improvements implemented in WGAN, we overcome common pitfalls faced in traditional GANs, achieving stable training and credible generative results. PyTorch, with its dynamic computation graph and extensive support for tensor operations, proves an invaluable asset in swiftly prototyping machine learning models like WGANs. By following this guide, you're well-equipped to explore more advanced generative models in deep learning.