<jupyter_start><jupyter_text>Pre-employment coursework Exercise 5 - Generative Adversarial Network- **Model:** A Vanilla GAN with WGAN-GP loss- **Aim:** Model should be trained to generate new instances of the training data.- **Sources:** Gulrajani, (2017), *Improved Training of Wasserstein GANs*, https://arxiv.org/pdf/1704.00028.pdf<jupyter_code># Imports
from collections import OrderedDict

import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.autograd import Variable, grad
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
torch.set_printoptions(linewidth=120)
torch.set_grad_enabled(True)

import torchvision
from torchvision.transforms import transforms

from IPython.display import display, clear_output

from utils_MNIST import *

print(torch.__version__)
print(torchvision.__version__)<jupyter_output><empty_output><jupyter_text>Load the MNIST dataset<jupyter_code># Define transform
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5,), (0.5,))])

# Load data
train_set = torchvision.datasets.MNIST(root='../data', train=True,
                                       download=True, transform=transform)

test_set = torchvision.datasets.MNIST(root='../data', train=False,
                                      download=True, transform=transform)<jupyter_output><empty_output><jupyter_text>Define model<jupyter_code>"""
class CriticNet(nn.Module):
"""
#A three hidden-layer critic neural network
"""
    def __init__(self):
        super(CriticNet, self).__init__()
        n_features = 784
        n_out = 1
        
        self.layers = nn.Sequential( 
            nn.Linear(n_features, 1024),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(1024, 512),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(256, n_out)
        )

    def forward(self, inputs):
        output = self.layers(inputs)
        return output
"""
    
class CriticNet(nn.Module):
    """
    A three hidden-layer critic neural network
    """
    def __init__(self):
        super(CriticNet, self).__init__()
        
        self.layers = nn.Sequential( 
            nn.Conv2d(1, 48, 3, 2, 1), # 48 14 14
            nn.LeakyReLU(0.1),
            nn.Conv2d(48, 96, 3, 2, 1), # 96 7 7
            nn.LeakyReLU(0.1),
            nn.Conv2d(96, 192, 3, 2, 1), # 192 4 4
            nn.LeakyReLU(0.1),
            nn.Flatten(start_dim=1),
            nn.Linear(192 * 4 * 4, 1)
        )

    def forward(self, inputs):
        output = self.layers(inputs)
        return output
class GeneratorNet(nn.Module):
    """
    A three hidden-layer generative neural network
    """
    def __init__(self, n_features):
        super(GeneratorNet, self).__init__()
        
        self.linear = nn.Linear(n_features, 4 * 4 * 192)
        self.BN = nn.BatchNorm1d(4 * 4 * 192)
        self.relu = nn.ReLU()
        self.layers = nn.Sequential(
            nn.ConvTranspose2d(192, 96, 2, 2, 0), # 96 8 8
            nn.BatchNorm2d(96),
            nn.ReLU(),
            nn.ConvTranspose2d(96, 48, 3, 2, 1), # 48 15 15
            nn.BatchNorm2d(48),
            nn.ReLU(),
            nn.ConvTranspose2d(48, 24, 2, 2, 1), # 24 28 28
            nn.BatchNorm2d(24),
            nn.ReLU(),
            nn.ConvTranspose2d(24, 1, 3, 1, 1), # 1 28 28
            nn.Tanh()
        )

    def forward(self, inputs):
        linear = self.relu(self.BN(self.linear(inputs)))
        linear = linear.view(-1, 192, 4, 4)
        generated_img = self.layers(linear)
        return generated_img
class WGAN_GP(nn.Module):
    """
    A Generative Adversarial Network model.
    """
    def __init__(self, latent_dim, lambd=10):
        super(WGAN_GP, self).__init__()
        self.latent_dim = latent_dim
        self.lambd = lambd
        self.critic = CriticNet()
        self.generator = GeneratorNet(self.latent_dim)
        self.d_optimizer = None
        self.g_optimizer = None
    
    @property
    def num_params(self):
        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
        num_p = sum([np.prod(p.size()) for p in model_parameters])
        return num_p
    
    def set_optimizers(self, d_optimizer, g_optimizer):
        self.d_optimizer = d_optimizer
        self.g_optimizer = g_optimizer
        
    def sample(self, batch_size):
        """
        Generates Gaussian sampled random values for image generation.
        """
        return Variable(torch.randn((batch_size, self.latent_dim)))
    
    def train_critic(self, data, train=True):
        """
        Trains critic model and returns WGAN loss and gradient penalty.
        """
        batch_size = data.size(0)
        samples = self.sample(batch_size)
        generated_data = self.generator(samples)
        gp = self.gradient_penalty(data, generated_data)
        d_loss = self.critic(generated_data).mean() - self.critic(data).mean() + gp
        self.d_optimizer.zero_grad()
        d_loss.backward()
        self.d_optimizer.step()
        return d_loss, gp
    
    def gradient_penalty(self, data, generated_data):
        """
        Calculates gradient penalty.
        """
        batch_size = data.size(0)
        #eps = torch.rand(batch_size, 1).expand_as(data)
        eps = torch.rand(batch_size, 1).expand_as(data.view((batch_size, -1)))
        eps = eps.view(batch_size, 1, 28, 28)
        
        interpolation = eps * data.data + (1 - eps) * generated_data.data
        interpolation = Variable(interpolation, requires_grad=True)

        interpolation_logits = self.critic(interpolation)
        grad_outputs = torch.ones(interpolation_logits.size())

        gradients = grad(outputs=interpolation_logits,
                         inputs=interpolation,
                         grad_outputs=grad_outputs,
                         create_graph=True,
                         retain_graph=True)[0]
        
        gradients = gradients.view(batch_size, -1)
        gp = self.lambd * ((gradients.norm(p=2, dim=1) - 1) ** 2).mean()
        return gp
    
    def train_generator(self, batch_size):
        """
        Trains generator model and returns generator loss.
        """
        samples = self.sample(batch_size)
        generated_data = self.generator(samples)
        g_loss = -self.critic(generated_data).mean()
        self.g_optimizer.zero_grad()
        g_loss.backward()
        self.g_optimizer.step()
        return g_loss
    
    def display(self, images, latents):
        """
        Displays real images and generated images.
        """
        print("Real images:")
        imgviz(images)
        
        with torch.no_grad():
            generated_images = self.generator(latents)
        print("Generated images:")
        imgviz(generated_images)

    def forward(self, inputs, train_generator=False):
        # Train critic
        data = Variable(inputs)
        d_loss, gp = self.train_critic(data)
        
        # If train_generator, train generator
        if train_generator:
            batch_size = data.size(0)
            g_loss = self.train_generator(batch_size)
            return d_loss, gp, g_loss
        else:
            return d_loss, gp
latent_dim = 128
print("Number of parameters in model with latent_dim={0}: {1}"
      .format(latent_dim, WGAN_GP(latent_dim).num_params))<jupyter_output>Number of parameters in model with latent_dim=128: 734162<jupyter_text>Training<jupyter_code>"""
parameters = OrderedDict(
    init_lr = [0.0001, 0.00005],
    batch_size = [128],
    latent_dim = [64],
    betas = [(0.5, 0.9), (0, 0.9), (0.9, 0.999)],
    lambd = [10, 5]
)
"""
param_names = ('init_lr', 'batch_size', 'latent_dim', 'betas', 'lambd')
parameters = OrderedDict(
    run1 = [0.0001, 128, 64, (0, 0.9), 10]
)
m = RunManager()
num_epochs = 150
num_samples = 20
num_critic_iters = 5

for hparams in RunBuilder.get_runs_from_params(param_names, parameters):

    # Instantiate a GAN model
    model = WGAN_GP(hparams.latent_dim, lambd=hparams.lambd)

    # Construct a DataLoader object with training data
    train_loader = DataLoader(train_set, batch_size=hparams.batch_size, shuffle=True)
    test_loader = DataLoader(test_set, batch_size=num_samples)
    test_images, _ = next(iter(test_loader))
    
    # Define random samples from standard Normal distribution for image generation
    test_latents = torch.randn((num_samples, hparams.latent_dim))
    
    # Define optimizers
    d_optimizer = optim.Adam(model.critic.parameters(), lr=hparams.init_lr, betas=hparams.betas)
    g_optimizer = optim.Adam(model.generator.parameters(), lr=hparams.init_lr, betas=hparams.betas)
    d_scheduler = optim.lr_scheduler.StepLR(d_optimizer, 60, 0.1)
    g_scheduler = optim.lr_scheduler.StepLR(g_optimizer, 60, 0.1)
    model.set_optimizers(d_optimizer, g_optimizer)
    
    # Setup run instance
    m.begin_run(hparams, model, test_images, test_latents)
    model.display(test_images, test_latents)
    print('Now training model with hyperparameters: init_lr={0}, batch_size={1}, latent_dim={2}, betas={3}, lambda={4}'
         .format(hparams.init_lr, hparams.batch_size, hparams.latent_dim, hparams.betas, hparams.lambd))
    
    # Start training loop
    for epoch in range(num_epochs):
        m.begin_epoch()
        
        # Train the model
        for i, batch in enumerate(train_loader):
            images, _ = batch

            # Calculating the loss
            if i % num_critic_iters == 0:
                d_loss, gp, g_loss = model(images, train_generator=True)
            else:
                d_loss, gp = model(images, train_generator=False)
                g_loss = None
            
            if i % 10 == 0:
                print('Epoch {0}, iteration {1}: d_loss {2}, gp {3}, g_loss {4}'
                      .format(epoch+1, i*hparams.batch_size,
                              round(d_loss.item(), 6),
                              round(gp.item(), 6),
                              round(g_loss.item(), 6)))

            m.track_loss(d_loss, gp, g_loss)
            
        m.end_epoch()
    
    #torch.save(model, './models/' + str(hparams) + '.pth')
    m.end_run()
    print("Model has finished training.\n")
    d_scheduler.step()
    g_scheduler.step()
    
m.save('results_MNIST_final')
print("Training completed.")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11456752655538695, 'fraction_numerical': 0.028641881638846737, 'mean_word_length': 2.3668582375478926, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 16, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '476384', 'license': 'MIT', 'n_tokens_mistral': 3647, 'n_tokens_neox': 3363, 'n_words': 724, 'path': 'gan/.ipynb_checkpoints/exercise5-model-wgangp-MNIST-checkpoint.ipynb', 'repo_name': 'jan-xu/autoencoders', 'size': '10544'}