Day 4 - Generative Adversarial Networks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generative Adversarial Networks\n",
    "So far, all the models we have worked with except the VAE have been discriminative models. This means that they are simply trying to predit something about our existing dataset. Sometimes, we would not like to discriminate but generate new examples of the data as in the case of video or image generation. Technically, this is the problem of modelling a probability distribution which we have samples of.\n",
    "\n",
    "One approach which implicitly models the distribution, the work of Ian Goodfellow, has be enjoying great success - often producing images indistinguishable from the examples on which it was trained. GANs take a game-theoretic approach, pitting two networks against eachother - the discriminator and the generator. The job of the generator is to produce images which are indistinguishable from the training set from latent variables while the job of the discriminator is to catch out the generator and discriminate real data from generated data. Initially they will both be terrible at their jobs but as the discriminator gets better, the generator is forced to get better to fool it and vice-versa. This loop continues until they are both excellent at their jobs and the generator can now be used to produce very realistic data points.\n",
    "\n",
    "![](GAN.png)\n",
    "\n",
    "An analogy often used to describe this is the detective and the forger. The generator is like a forger who is trying to produce paintings indistinguishable from other famous paintings by an artists while the discriminator is like a detective who is trying to catch the forger out. As the detective gets better at catching the generator out, the generator is forced to improve to fool the detective."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementation\n",
    "We will be training a GAN on the fashion MNIST dataset so we will be able to produce images of items of clothing which look like they came from the original dataset.\n",
    "\n",
    "We begin by importing the appropriate libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "import torch\n",
    "from torch.autograd import Variable\n",
    "import torchvision\n",
    "from torchvision import transforms, datasets\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We load our dataset into a pytorch dataloader which we will use later to produce random batches of samples from our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 100\n",
    "\n",
    "train_data = datasets.FashionMNIST(root='fashiondata/',\n",
    "                                 transform=transforms.ToTensor(),\n",
    "                                 train=True,\n",
    "                                 download=True\n",
    "                                 )\n",
    "\n",
    "train_samples = torch.utils.data.DataLoader(dataset=train_data,\n",
    "                                           batch_size=batch_size,\n",
    "                                           shuffle=True\n",
    "                                            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the NN model we will be using for our discriminator. It takes in the 28x28 image and performs convolutions followed by one fully connected layer to output the probability of the data point being real and not generated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class discriminator(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = torch.nn.Conv2d(1, 64, kernel_size=4, stride=2, padding=1) #1x28x28-> 64x14x14\n",
    "        self.conv2 = torch.nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1) #64x14x14-> 128x7x7\n",
    "        self.dense1 = torch.nn.Linear(128*7*7, 1)\n",
    "\n",
    "        self.bn1 = torch.nn.BatchNorm2d(64)\n",
    "        self.bn2 = torch.nn.BatchNorm2d(128)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.bn1(self.conv1(x)))\n",
    "        x = F.relu(self.bn2(self.conv2(x))).view(-1, 128*7*7)\n",
    "        x = F.sigmoid(self.dense1(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the NN model for the generator. This takes in a latent vector of size 128 and performs fully connected layers followed by upconvolution to output us a 28x28 image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class generator(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.dense1 = torch.nn.Linear(128, 256)\n",
    "        self.dense2 = torch.nn.Linear(256, 1024)\n",
    "        self.dense3 = torch.nn.Linear(1024, 128*7*7)\n",
    "        self.uconv1 = torch.nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1) #128x7x7 -> 64x14x14\n",
    "        self.uconv2 = torch.nn.ConvTranspose2d(64, 1, kernel_size=4, stride=2, padding=1) #64x14x14 -> 1x28x28\n",
    "\n",
    "        self.bn1 = torch.nn.BatchNorm1d(256)\n",
    "        self.bn2 = torch.nn.BatchNorm1d(1024)\n",
    "        self.bn3 = torch.nn.BatchNorm1d(128*7*7)\n",
    "        self.bn4 = torch.nn.BatchNorm2d(64)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.bn1(self.dense1(x)))\n",
    "        x = F.relu(self.bn2(self.dense2(x)))\n",
    "        x = F.relu(self.bn3(self.dense3(x))).view(-1, 128, 7, 7)\n",
    "        x = F.relu(self.bn4(self.uconv1(x)))\n",
    "        x = F.sigmoid(self.uconv2(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We instantiate our models from the classes we define and their optimizers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#instantiate model\n",
    "d = discriminator()\n",
    "g = generator()\n",
    "\n",
    "#training hyperparameters\n",
    "no_epochs = 100\n",
    "dlr = 0.0003\n",
    "glr = 0.0003\n",
    "\n",
    "d_optimizer = torch.optim.Adam(d.parameters(), lr=dlr)\n",
    "g_optimizer = torch.optim.Adam(g.parameters(), lr=glr)\n",
    "\n",
    "dcosts = []\n",
    "gcosts = []\n",
    "plt.ion()\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "ax.set_xlabel('Epoch')\n",
    "ax.set_ylabel('Cost')\n",
    "ax.set_xlim(0, no_epochs)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define the training loop. For every batch of training data that we look at, we get the generator to produce an equally sized batch of generated images. We then get the discriminator to make predictions on both sets and calulate the cost for both networks before calculating the gradients and training each one in turn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(no_epochs, glr, dlr):\n",
    "    for epoch in range(no_epochs):\n",
    "        epochdcost = 0\n",
    "        epochgcost = 0\n",
    "\n",
    "        #iteratre over mini-batches\n",
    "        for k, (real_images, _ ) in enumerate(train_samples):\n",
    "            real_images = Variable(real_images) #real images from training set\n",
    "\n",
    "            z = Variable(torch.randn(batch_size, 128)) #generate random latent variable to generate images\n",
    "            generated_images = g.forward(z) #generate images\n",
    "\n",
    "            gen_pred = d.forward(generated_images) #prediction of generator on generated batch\n",
    "            real_pred = d.forward(real_images) #prediction of generator on real batch\n",
    "\n",
    "            dcost = -torch.sum(torch.log(real_pred)) - torch.sum(torch.log(1-gen_pred)) #cost of discriminator\n",
    "            gcost = -torch.sum(torch.log(gen_pred))/batch_size #cost of generator\n",
    "            \n",
    "            #train discriminator\n",
    "            d_optimizer.zero_grad()\n",
    "            dcost.backward(retain_graph=True) #retain the computation graph so we can train generator after\n",
    "            d_optimizer.step()\n",
    "            \n",
    "            #train generator\n",
    "            g_optimizer.zero_grad()\n",
    "            gcost.backward()\n",
    "            g_optimizer.step()\n",
    "\n",
    "            epochdcost += dcost.data[0]\n",
    "            epochgcost += gcost.data[0]\n",
    "            \n",
    "            #give us an example of a generated image after every 10000 images produced\n",
    "            if k*batch_size%10000 ==0:\n",
    "                g.eval() #put in evaluation mode\n",
    "                noise_input = Variable(torch.randn(1, 128))\n",
    "                generated_image = g.forward(noise_input)\n",
    "\n",
    "                plt.figure(figsize=(1, 1))\n",
    "                plt.imshow(generated_image.data[0][0], cmap='gray_r')\n",
    "                plt.show()\n",
    "                g.train() #put back into training mode\n",
    "\n",
    "\n",
    "            epochdcost /= 60000/batch_size\n",
    "            epochgcost /= 60000/batch_size\n",
    "\n",
    "            print('Epoch: ', epoch)\n",
    "            print('Disciminator cost: ', epochdcost)\n",
    "            print('Generator cost: ', epochgcost)\n",
    "        \n",
    "        #plot costs\n",
    "        \n",
    "        dcosts.append(epochdcost)\n",
    "        gcosts.append(epochgcost)\n",
    "\n",
    "        ax.plot(dcosts, 'b')\n",
    "        ax.plot(gcosts, 'r')\n",
    "\n",
    "        fig.canvas.draw()\n",
    "\n",
    "train(no_epochs, glr, dlr)"
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}