lsgan2.py

from __future__ import print_function
from __future__ import division

import argparse

import matplotlib.pyplot as plot
import numpy as np
import torch.utils.data
import torchvision.datasets
from torch import nn, optim
from torch.autograd import Variable
from torchvision import transforms


class Reshape(nn.Module):
    def __init__(self, *shape):
        super(Reshape, self).__init__()
        self.shape = tuple(shape)

    def forward(self, input):
        return input.view(len(input), *self.shape)

    def __repr__(self):
        return '{}{}'.format(self.__class__.__name__, self.shape)


class Flatten(nn.Module):
    def __init__(self):
        super(Flatten, self).__init__()

    def forward(self, input):
        return input.view(len(input), int(np.prod(input.shape[1:])))


class AddNoise(nn.Module):
    def __init__(self, mean=0.0, stddev=0.1):
        super(AddNoise, self).__init__()
        self.mean = mean
        self.stddev = stddev

    def forward(self, images):
        return images + torch.empty_like(images).normal_()


parser = argparse.ArgumentParser()
parser.add_argument('-n', '--noise-size', type=int, default=100)
parser.add_argument('-e', '--epochs', type=int, default=30)
parser.add_argument('-b', '--batch-size', type=int, default=256)
parser.add_argument('-x', '--examples', type=int, default=8)
parser.add_argument('-s', '--show', action='store_true')
parser.add_argument('-g', '--gpu', action='store_true')
options = parser.parse_args()


def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)


class Generator(nn.Module):
    def __init__(self, noise_size):
        super(Generator, self).__init__()

        self.model = nn.Sequential(
            # Layer 1
            nn.Linear(noise_size, 7 * 7 * 256),
            nn.BatchNorm1d(7 * 7 * 256, momentum=0.9),
            nn.LeakyReLU(0.2, inplace=True),
            Reshape(256, 7, 7),
            # Layer 2
            nn.Upsample(scale_factor=2),
            nn.Conv2d(256, 128, kernel_size=5, padding=2, bias=False),
            nn.BatchNorm2d(128, momentum=0.9),
            nn.LeakyReLU(0.2, inplace=True),
            # Layer 3
            nn.Upsample(scale_factor=2),
            nn.Conv2d(128, 64, kernel_size=5, padding=2, bias=False),
            nn.BatchNorm2d(64, momentum=0.9),
            nn.LeakyReLU(0.2, inplace=True),
            # Layer 4
            nn.Conv2d(64, 32, kernel_size=5, padding=2, bias=False),
            nn.BatchNorm2d(32, momentum=0.9),
            nn.LeakyReLU(0.2, inplace=True),
            # Output
            nn.Conv2d(32, 1, kernel_size=5, padding=2, bias=False),
            nn.Tanh())

    def forward(self, noise):
        return self.model(noise.squeeze())


class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.model = nn.Sequential(
            AddNoise(),
            # Layer 1
            nn.Conv2d(1, 32, kernel_size=5, padding=2, stride=2, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # Layer 2
            nn.Conv2d(32, 64, kernel_size=5, padding=2, stride=2, bias=False),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2, inplace=True),
            # Layer 3
            nn.Conv2d(64, 128, kernel_size=5, padding=2, stride=2, bias=False),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            # Layer 4
            nn.Conv2d(128, 256, kernel_size=5, padding=2, bias=False),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            # Output
            Flatten(),
            nn.Linear(256 * 4 * 4, 1))

    def forward(self, images):
        return self.model(images)


generator = Generator(options.noise_size)
generator.apply(weights_init)
print(generator)

discriminator = Discriminator()
discriminator.apply(weights_init)
print(discriminator)

dataset = torchvision.datasets.MNIST(
    root='./data',
    train=True,
    download=True,
    transform=transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize([0.5], [0.5])]))
data_loader = torch.utils.data.DataLoader(
    dataset,
    batch_size=options.batch_size,
    num_workers=4,
    shuffle=True,
    drop_last=True)

criterion = nn.MSELoss()

discriminator_optimizer = optim.Adam(
    discriminator.parameters(), lr=2e-4, betas=(0.5, 0.999))
generator_optimizer = optim.Adam(
    generator.parameters(), lr=2e-4, betas=(0.5, 0.999))

noise = Variable(torch.Tensor(options.batch_size, options.noise_size))
real_labels = Variable(torch.Tensor(options.batch_size))
fake_labels = Variable(torch.zeros(options.batch_size))

if options.gpu:
    print('Copying data to GPU memory ...')
    generator.cuda()
    discriminator.cuda()
    criterion.cuda()
    noise = noise.cuda()
    real_labels = real_labels.cuda()
    fake_labels = fake_labels.cuda()

try:
    for epoch_index in range(1, options.epochs + 1):
        for batch_index, batch in enumerate(data_loader):
            real_images = Variable(batch[0])
            if options.gpu:
                real_images = real_images.cuda()

            discriminator.zero_grad()
            real_predictions = discriminator(real_images)
            real_labels.uniform_(0.8, 1.0)
            d_loss_real = criterion(real_predictions, real_labels)
            d_loss_real.backward()

            noise.normal_()
            fake_images = generator(noise)
            fake_predictions = discriminator(fake_images.detach())
            fake_labels.fill_(0)
            d_loss_fake = criterion(fake_predictions, fake_labels)
            d_loss_fake.backward()

            d_loss_value = d_loss_real.data.mean() + d_loss_fake.data.mean()
            discriminator_optimizer.step()

            generator.zero_grad()
            fake_predictions = discriminator(fake_images)
            fake_labels.fill_(1)
            g_loss = criterion(fake_predictions, fake_labels)
            g_loss.backward()
            g_loss_value = g_loss.data.mean()
            generator_optimizer.step()

            message = 'Epoch: {}/{} | Batch: {}/{} | G: {:.10f} D: {:.10f}'
            message = message.format(epoch_index, options.epochs, batch_index,
                                     len(data_loader), g_loss_value,
                                     d_loss_value)
            print('\r{}'.format(message), end='')
except KeyboardInterrupt:
    pass

print('\nTraining complete!')

if not options.show:
    plot.switch_backend('Agg')

noise = Variable(torch.Tensor(options.examples, options.noise_size).normal_())
if options.gpu:
    noise = noise.cuda()
images = generator(noise).detach().cpu().numpy()
images = (images + 1) / 2
print(images[0])
print(images.mean(), images.std())

number_of_columns = 4
number_of_rows = len(images) // number_of_columns
plot.figure()
for i, image in enumerate(images):
    plot.subplot(number_of_rows, number_of_columns, 1 + i)
    plot.axis('off')
    plot.imshow(image.squeeze(), cmap='gray')

if options.show:
    print('Showing plots')
    plot.show()
else:
    print('Saving figure.png')
    plot.savefig('figure.png')