<filename>basic_gan.py
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader

class Generator(nn.Module):
    
    def __init__(self):
        super(Generator, self).__init__()
        self.fc1 = nn.Linear(256, 384)
        self.bn1 = nn.BatchNorm1d(384)
        self.fc2 = nn.Linear(384, 512)
        self.bn2 = nn.BatchNorm1d(512)
        self.fc3 = nn.Linear(512, 784)
    
    def forward(self, z):
        out = nn.functional.leaky_relu(self.bn1(self.fc1(z)), 0.05)
        out = nn.functional.leaky_relu(self.bn2(self.fc2(out)), 0.05)
        out = torch.sigmoid(self.fc3(out))
        return out

class Critic(nn.Module):
    
    def __init__(self):
        super(Critic, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.bn1 = nn.BatchNorm1d(256)
        self.fc2 = nn.Linear(256, 128)
        self.bn2 = nn.BatchNorm1d(128)
        self.fc3 = nn.Linear(128, 1)
    
    def forward(self, x):
        out = nn.functional.leaky_relu(self.bn1(self.fc1(x)), 0.05)
        out = nn.functional.leaky_relu(self.bn2(self.fc2(out)), 0.05)
        out = self.fc3(out)
        return out

class GAN(object):
    
    def __init__(self, lr, batch_size, device):
        self.batch_size = batch_size
        self.device = device
        self.generator = Generator().to(self.device)
        self.critic = Critic().to(self.device)
        self.g_optimiser = torch.optim.Adam(self.generator.parameters(), lr=lr, betas=(0.5, 0.99))
        self.c_optimiser = torch.optim.Adam(self.critic.parameters(), lr=lr, betas=(0.5, 0.99))
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Lambda(lambda x: x.view(-1))
        ])
        dataset = MNIST(".", train=True, download=True, transform=transform)
        self.dataloader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True)
        self.loss_func = nn.BCEWithLogitsLoss()
    
    def sample_z(self, n):
        z = (torch.rand(n, 256)).to(self.device)
        return z
    
    def generate_images(self, z):
        self.generator.eval()
        imgs = self.generator(z)
        imgs = imgs.view(imgs.size(0), 28, 28).detach().cpu().data.numpy()
        return imgs
    
    def fit_networks(self, iterations, plot_every):
        for i in range(iterations):
            self.critic.train()
            self.generator.eval()
            self.c_optimiser.zero_grad()
            z = self.sample_z(self.batch_size)
            for data, labels in self.dataloader:
                x = data.to(self.device)
                break
            x_hat = self.generator.forward(z)
            y_hat = self.critic.forward(torch.cat((x, x_hat), dim=0))
            labels = torch.cat((torch.ones(self.batch_size, 1),
                                torch.zeros(self.batch_size, 1)),
                               dim=0).to(self.device)
            critic_loss = self.loss_func(y_hat, labels)
            critic_loss.backward()
            self.c_optimiser.step()
            
            self.critic.eval()
            self.generator.train()
            self.g_optimiser.zero_grad()
            z = self.sample_z(self.batch_size)
            x_hat = self.generator.forward(z)
            y_hat = self.critic.forward(x_hat)
            labels = torch.ones(self.batch_size, 1).to(self.device)
            generator_loss = self.loss_func(y_hat, labels)
            generator_loss.backward()
            self.g_optimiser.step()
            
            if (i+1) % plot_every == 0:
                z = self.sample_z(4)
                imgs = self.generate_images(z)
                fig, ax = plt.subplots(figsize=(12, 3), nrows=1, ncols=4)
                for j in range(4):
                    ax[j].imshow(imgs[j], cmap='gray')
                    ax[j].set_title(str(i+1) + ' iterations')
                plt.show()

lr = 0.001
batch_size = 256
device = 'cuda' if torch.cuda.is_available() else 'cpu'

iterations = 20000
plot_every = 200

gan = GAN(lr, batch_size, device)
gan.fit_networks(iterations, plot_every)