GAN-Notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generative Adversarial Network - PyTorch\n",
    "\n",
    "This notebook contains the same implementation of a normal linear GAN as the python files. The notebook is primarily meant to give a better understanding of the model without the pain to search through the python files. In addition to that code for a run in collaboratory is provided.\n",
    "\n",
    "\n",
    "GANs were [first introduced](https://arxiv.org/abs/1406.2661) in 2014 from Ian Goodfellow and others. \n",
    "GANs consists of two networks: a generator ***G*** and a discriminator ***D***. The generator is designed to generate realistic looking data \"fake data\" and the discriminator is designed to distinguish between generated \"fake data\" and real data from a dataset. Both networks play a game against each other where the generator tries to fool the discriminator by generating data as real as possible and the discriminator tries to classify fake and real data correctly.\n",
    "\n",
    "* The generator is a linear feedforward network which takes a random vector as input and outputs another vector which can be reshaped to an image\n",
    "* The discriminator is a linear classifier that takes an image vector as input and outputs a probability of that image being real.\n",
    "\n",
    "This game leads to a generator producing data that is indistinguishable from real data to the discriminator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Google Collaboratory code\n",
    "Uncommend the following cells if you're running this notebook in google collab. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#-----MOUNT DRIVE-----\n",
    "#from google.colab import drive\n",
    "#drive.mount('/gdrive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#-----INSTALL PYTORCH-----\n",
    "#from os.path import exists\n",
    "#from wheel.pep425tags import get_abbr_impl, get_impl_ver, get_abi_tag\n",
    "#platform = '{}{}-{}'.format(get_abbr_impl(), get_impl_ver(), get_abi_tag())\n",
    "#cuda_output = !ldconfig -p|grep cudart.so|sed -e 's/.*\\.\\([0-9]*\\)\\.\\([0-9]*\\)$/cu\\1\\2/'\n",
    "#accelerator = cuda_output[0] if exists('/dev/nvidia0') else 'cpu'\n",
    "#!pip install -q http://download.pytorch.org/whl/{accelerator}/torch-0.4.1-{platform}-linux_x86_64.whl torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#-----REINSTALL PILLOW-----\n",
    "#!pip install Pillow==4.0.0\n",
    "#!pip install PIL\n",
    "#!pip install image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# Import PyTorch dependencies\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "import torch.nn.functional as F\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make directories for images and mnist dataset\n",
    "os.makedirs('mnist_data', exist_ok=True)\n",
    "os.makedirs('generated_images', exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import and Transform Data\n",
    "[MNIST](http://yann.lecun.com/exdb/mnist/) will be used as an example dataset. For a custom dataset change the *custom_image_path* parameter. A grayscale image is required."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#-----DATASET PARAMETERS-----\n",
    "custom_image_path = None\n",
    "batch_size = 128\n",
    "image_size = 28\n",
    "#---------------------------\n",
    "\n",
    "# Resize data\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize(image_size),\n",
    "    transforms.Grayscale(),\n",
    "    transforms.ToTensor()\n",
    "])\n",
    "\n",
    "if custom_image_path != None:\n",
    "    train_dataset = datasets.ImageFolder(custom_image_path)\n",
    "else:\n",
    "    train_dataset = datasets.MNIST(root=\"mnist_data\", train=True, download=True, transform=transform)\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def im_convert(tensor, rescale=False):\n",
    "    tensor = tensor.view(1, image_size, image_size)\n",
    "    image = tensor.to(\"cpu\").clone().detach()\n",
    "    image = image.numpy().squeeze()\n",
    "    if rescale:\n",
    "        # Rescale image from tanh output (1, -1) to rgb (0, 255)\n",
    "        image = ((image + 1) * 255 / (2)).astype(np.uint8)\n",
    "    return image\n",
    "\n",
    "# Show one image \n",
    "dataiter = iter(train_loader)\n",
    "images, labels = dataiter.next()\n",
    "image = im_convert(images[0])\n",
    "plt.imshow(image, cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Scale image values from -1 to 1 to be close to the output of the tanh function\n",
    "def scale(x, feature_range=(-1, 1)):\n",
    "    min, max = feature_range\n",
    "    x = x * (max-min) + min\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GAN Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Discriminator\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, image_size):\n",
    "        super().__init__()\n",
    "        input_features = 1 * image_size * image_size\n",
    "    \n",
    "        # Define hidden linear layers\n",
    "        self.hcl1 = nn.Linear(input_features, 1024)\n",
    "        self.hcl2 = nn.Linear(1024, 512)\n",
    "        self.hcl3 = nn.Linear(512, 256)\n",
    "        self.out = nn.Linear(256, 1)\n",
    "        # Activation function\n",
    "        self.leaky_relu = nn.LeakyReLU(negative_slope=0.2)\n",
    "  \n",
    "    def forward(self, x):\n",
    "        x = self.leaky_relu(self.hcl1(x))\n",
    "        x = self.leaky_relu(self.hcl2(x))\n",
    "        x = self.leaky_relu(self.hcl3(x))\n",
    "        x = torch.sigmoid(self.out(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def linear_block(in_features, out_features, batch_norm=True):\n",
    "    layers = []\n",
    "    if batch_norm:\n",
    "        linear_layer = nn.Linear(in_features, out_features, bias=False)\n",
    "        batch_norm = nn.BatchNorm1d(out_features)\n",
    "        layers = [linear_layer, batch_norm]\n",
    "    else:\n",
    "        layers.append(nn.Linear(in_features, out_features))\n",
    "    return nn.Sequential(*layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "    def __init__(self, input_features, image_size):\n",
    "        super().__init__()\n",
    "        output_features = 1 * image_size * image_size\n",
    "    \n",
    "        # Define hidden linear layers\n",
    "        self.hcl1 = linear_block(input_features, 256)\n",
    "        self.hcl2 = linear_block(256, 512)\n",
    "        self.hcl3 = linear_block(512, 1024)\n",
    "        self.output = linear_block(1024, output_features)\n",
    "    \n",
    "        # Activation function\n",
    "        self.leaky_relu = nn.LeakyReLU(negative_slope=0.2)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.leaky_relu(self.hcl1(x))\n",
    "        x = self.leaky_relu(self.hcl2(x))\n",
    "        x = self.leaky_relu(self.hcl3(x))\n",
    "        x = torch.tanh(self.output(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build the models\n",
    "z_size = 100\n",
    "\n",
    "D = Discriminator(image_size=image_size)\n",
    "G = Generator(input_features=z_size, image_size=image_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If a gpu is available move all models to gpu\n",
    "if torch.cuda.is_available():\n",
    "    G = G.cuda()\n",
    "    D = D.cuda()\n",
    "    print(\"GPU available. Moved models to GPU.\")\n",
    "else:\n",
    "    print(\"Training on CPU.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def real_loss(predictions, smooth=False):\n",
    "    batch_size = predictions.shape[0]\n",
    "    labels = torch.ones(batch_size)\n",
    "    # Smooth labels for discriminator to weaken learning\n",
    "    if smooth:\n",
    "        labels = labels * 0.9\n",
    "    # We use the binary cross entropy loss | Model has a sigmoid function\n",
    "    criterion = nn.BCELoss()\n",
    "    # Move models to GPU if available\n",
    "    if torch.cuda.is_available():\n",
    "        labels = labels.cuda()\n",
    "        criterion = criterion.cuda()\n",
    "    loss = criterion(predictions.squeeze(), labels)\n",
    "    return loss\n",
    "\n",
    "def fake_loss(predictions):\n",
    "    batch_size = predictions.shape[0]\n",
    "    labels = torch.zeros(batch_size)\n",
    "    criterion = nn.BCELoss()\n",
    "    # Move models to GPU if available\n",
    "    if torch.cuda.is_available():\n",
    "        labels = labels.cuda()\n",
    "        criterion = criterion.cuda()\n",
    "    loss = criterion(predictions.squeeze(), labels)\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def random_vector(batch_size, length):\n",
    "    # Sample from a Gaussian distribution\n",
    "    z_vec = torch.randn(batch_size, length).float()\n",
    "    if torch.cuda.is_available():\n",
    "        z_vec = z_vec.cuda()\n",
    "    return z_vec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#-----TRAINING PARAMETERS-----\n",
    "lr = 0.0002\n",
    "beta1 = 0.5\n",
    "beta2 = 0.999\n",
    "#-----------------------------\n",
    "\n",
    "d_optimizer = optim.Adam(D.parameters(), lr=lr, betas=[beta1, beta2])\n",
    "g_optimizer = optim.Adam(G.parameters(), lr=lr, betas=[beta1, beta2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_discriminator(generator, discriminator, optimizer, real_data, batch_size, z_size):\n",
    "    # Reshape real_data to vector\n",
    "    real_data = real_data.view(batch_size, -1)\n",
    "    # Rescale real_data to range -1 - 1\n",
    "    real_data = scale(real_data)\n",
    "    \n",
    "    # Reset gradients and set model to training mode\n",
    "    optimizer.zero_grad()\n",
    "    discriminator.train()\n",
    "    \n",
    "    # Train on real data\n",
    "    real_data_logits = discriminator.forward(real_data)\n",
    "    loss_real = real_loss(real_data_logits, smooth=True)\n",
    "    # Generate fake data\n",
    "    z_vec = random_vector(batch_size, z_size)\n",
    "    fake_data = generator.forward(z_vec)\n",
    "    # Train on fake data\n",
    "    fake_data_logits = discriminator.forward(fake_data)\n",
    "    loss_fake = fake_loss(fake_data_logits)\n",
    "    # Calculate total loss\n",
    "    total_loss = loss_real + loss_fake\n",
    "    total_loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    return total_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_generator(generator, discriminator, optimizer, batch_size, z_size):\n",
    "    # Reset gradients and set model to training mode\n",
    "    optimizer.zero_grad()\n",
    "    generator.train()\n",
    "    # Generate fake data\n",
    "    z_vec = random_vector(batch_size, z_size)\n",
    "    fake_data = generator.forward(z_vec)\n",
    "    # Train generator with output of discriminator\n",
    "    discriminator_logits = discriminator.forward(fake_data)\n",
    "    # Reverse labels\n",
    "    loss = real_loss(discriminator_logits)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Trainings loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 50\n",
    "# After how many batches should generated sample images be shown?\n",
    "print_every = 200\n",
    "# How many images should be shown?\n",
    "sample_size = 8\n",
    "# After how many epochs should the loss be plotted?\n",
    "plot_every = 5\n",
    "# Create some sample noise\n",
    "sample_noise = random_vector(sample_size, z_size)\n",
    "#-------------------------\n",
    "\n",
    "# Keep track of losses\n",
    "d_losses = []\n",
    "g_losses = []\n",
    "\n",
    "for e in range(epochs):\n",
    "    for batch_i, (images, _) in enumerate(train_loader):\n",
    "        batch_size = images.shape[0]\n",
    "        # Move images to GPU if available\n",
    "        if torch.cuda.is_available():\n",
    "            images = images.cuda()\n",
    "        # Train discriminator\n",
    "        d_loss = train_discriminator(G, D, d_optimizer, images, batch_size, z_size)\n",
    "        # Train generator\n",
    "        g_loss = train_generator(G, D, g_optimizer, batch_size, z_size)\n",
    "        \n",
    "        # Keep track of losses\n",
    "        d_losses.append(d_loss)\n",
    "        g_losses.append(g_loss)\n",
    "        \n",
    "        # Print some sample pictures\n",
    "        if (batch_i % print_every == 0):\n",
    "            print(\"Epoch: {}, Batch: {}, D-Loss: {}, G-Loss: {}\".format(e, batch_i, d_loss, g_loss))\n",
    "            # Make sample generation\n",
    "            G.eval()\n",
    "            fig, axes = plt.subplots(1, sample_size, figsize=(20, 10))\n",
    "            # Generate predictions\n",
    "            predictions = G.forward(sample_noise)\n",
    "            for i in range(sample_size):\n",
    "                axes[i].imshow(im_convert(predictions[i], rescale=True), cmap=\"gray\")\n",
    "            plt.show()\n",
    "    if (e % plot_every == 0):\n",
    "        # Print losses\n",
    "        plt.plot(d_losses, label=\"Discriminator\", alpha=0.5)\n",
    "        plt.plot(g_losses, label=\"Generator\", alpha=0.5)\n",
    "        plt.title(\"Trainings loss\")\n",
    "        plt.legend()\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(d_losses, label=\"Discriminator\", alpha=0.5)\n",
    "plt.plot(g_losses, label=\"Generator\", alpha=0.5)\n",
    "plt.title(\"Trainings loss\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_samples(num_samples):\n",
    "    G.eval()\n",
    "    z_vec = random_vector(num_samples, z_size)\n",
    "    predictions = G.forward(z_vec)\n",
    "    fig, axes = plt.subplots(1, num_samples, figsize=(20, 10))\n",
    "    for i in range(num_samples):\n",
    "        axes[i].imshow(im_convert(predictions[i], rescale=True), cmap=\"gray\")              \n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_samples(8)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}