wgan-gp's Introduction

Wasserstein GAN with Gradient penalty

Pytorch implementation of Improved Training of Wasserstein GANs by Gulrajani et al.

Examples

MNIST

Parameters used were lr=1e-4, betas=(.9, .99), dim=16, latent_dim=100. Note that the images were resized from (28, 28) to (32, 32).

Training (200 epochs)

Samples

Fashion MNIST

Training (200 epochs)

Samples

LSUN Bedrooms

Gif [work in progress]

Samples [work in progress]

Usage

Set up a generator and discriminator model

from models import Generator, Discriminator
generator = Generator(img_size=(32, 32, 1), latent_dim=100, dim=16)
discriminator = Discriminator(img_size=(32, 32, 1), dim=16)

The generator and discriminator are built to automatically scale with image sizes, so you can easily use images from your own dataset.

Train the generator and discriminator with the WGAN-GP loss

import torch
# Initialize optimizers
G_optimizer = torch.optim.Adam(generator.parameters(), lr=1e-4, betas=(.9, .99))
D_optimizer = torch.optim.Adam(discriminator.parameters(), lr=1e-4, betas=(.9, .99))

# Set up trainer
from training import Trainer
trainer = Trainer(generator, discriminator, G_optimizer, D_optimizer,
                  use_cuda=torch.cuda.is_available())

# Train model for 200 epochs
trainer.train(data_loader, epochs=200, save_training_gif=True)

This will train the models and generate a gif of the training progress.

Note that WGAN-GPs take a long time to converge. Even on MNIST it takes about 50 epochs to start seeing decent results. For more information and a full example on MNIST, check out main.py.