pegasus.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pegasus"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Google COLAB Settings\n",
    "In this section are certain processes that should be run when running the code through Google COLAB so to have access to a GPU. If such is the case, uncomment the sections and run them sequentially. otherwise, feel free to skip directly to [Imports](#Imports)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Installs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "X0XeNJMELfIb"
   },
   "outputs": [],
   "source": [
    "# %%capture\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\n",
    "# !pip install livelossplot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Google Drive\n",
    "This portion is exclusively for development on _my_ end. I use Google Drive to access the training/testing data without having to redownload it each time the Google COLAB runtime is reset. \n",
    "\n",
    "Of course anyone who does not have access to my Google credentials will not be able to access my Drive. As such, these users should skip directly to [Imports](#Imports). The result will be that torchvision will personally download the CIFAR data from the web each time the COLAB runtime is reset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Mounting Drive\n",
    "This mounts Google Drive to the local runtime. If Drive is already mounted, then of course, it will not try to mount it again. It will of course ask for authentication."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from google.colab import drive\n",
    "# drive.mount('/content/gdrive')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Importing data from Drive\n",
    "Here I import the CIFAR data which I have previously downloaded and stored in my Google Drive. To do so I copy the corresponding directory from my Google Drive into the COLAB Runtime to avoid having to redownload it each time my COLAB runtime is reset. The ```-n``` flag is set to avoid overwriting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !cp -r -n /content/gdrive/My\\ Drive/Education/Undergraduate/Year_3/Computer_Science/SSA/Machine_Learning/Coursework/ML_Classifier-Pegasus-Generator/data/ /content/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6N22Uz-kLiZW"
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MK1Jl7nkLnPA"
   },
   "outputs": [],
   "source": [
    "from time import sleep\n",
    "import math\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.autograd import Variable\n",
    "import torch.nn.functional as F\n",
    "import torchvision\n",
    "import matplotlib.pyplot as plt\n",
    "from livelossplot import PlotLosses\n",
    "\n",
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "print(f'Device being used: {device}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "run1dh_hM0oO"
   },
   "source": [
    "## Dataset Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 86
    },
    "colab_type": "code",
    "id": "bK383zeDM4Ac",
    "outputId": "5a7bb37e-9b16-4f87-9c7e-6b07c5ed8d2e"
   },
   "outputs": [],
   "source": [
    "# define class names for CIFAR 10\n",
    "class_names = ['airplane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define transforms to be applied to training\n",
    "transforms = torchvision.transforms.Compose([\n",
    "        torchvision.transforms.ToTensor(),\n",
    "        torchvision.transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n",
    "])\n",
    "\n",
    "inverse_normalize = torchvision.transforms.Normalize(\n",
    "    mean=[-1, -1, -1],\n",
    "    std=[2, 2, 2]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Getting the Data\n",
    "If not already present, this cell will download the [CIFAR 10](https://www.cs.toronto.edu/~kriz/cifar.html) dataset from the web. Otherwise it will simply read it from the existing directory. The transforms defined [above](#Transforms) will be applied.\n",
    "\n",
    "#### Retaining only Relevant Data\n",
    "It should be noted that the eventual goal of this script is to generate a pegasus -- that is, a horse with wings. As such, there is no need for training the network on data that is not relevant. The network might as well only focus on creating inner representations of horses and wings. Therefore it makes more sense to only train the network on images of horses, birds and planes (after all no one said this _couldn't_ be a cyborg-pegasus). It is hypothesized that this will make training more efficient and will render [mode collapse](https://arxiv.org/pdf/1611.02163.pdf) less of an issue, as there are less modes to collapse to.\n",
    "\n",
    "To achieve this effect, we make use of the [SubsetRandomSampler class provided by torch.utils.data](https://pytorch.org/docs/stable/data.html#torch.utils.data.SubsetRandomSampler) passing to it a list of the indices that correspond exclusively to these classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download training set\n",
    "train_set = torchvision.datasets.CIFAR10('data', train=True, download=True, transform=transforms)\n",
    "\n",
    "# get the labels (numbers) corresponding to airplane, horse, bird class names\n",
    "accepted_labels = [i for i in range(len(class_names)) if class_names[i] in ['airplane', 'horse', 'bird']]\n",
    "\n",
    "# get the indices in the downloaded data sets corresponding to airplanes, horses and birds.\n",
    "relevant_train_indices = [i for i in range(len(train_set.targets)) if train_set.targets[i] in accepted_labels]\n",
    "\n",
    "# insatiating the samplers to feed into the DataLoader so that only airplanes, horses and birds are loaded\n",
    "train_sampler = torch.utils.data.SubsetRandomSampler(relevant_train_indices)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading the Data\n",
    "Having obtained the data, it needs to be loaded into an iterable format for pytorch to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define batch size\n",
    "BATCH_SIZE = 128\n",
    "\n",
    "# load training set into a data torch data object, shuffled\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size=BATCH_SIZE, drop_last=True, sampler = train_sampler)\n",
    "\n",
    "# diagnostic print size of training dataset\n",
    "print(f'> Size of training dataset: {len(relevant_train_indices)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "q-FdW5HnimG2"
   },
   "source": [
    "### Viewing (some of) the Data\n",
    "...It almost seems as if birds are not worth training on since their wings are often closed..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 589
    },
    "colab_type": "code",
    "id": "BtJs-qxHRLXz",
    "outputId": "c504705d-1713-4251-8698-6a260807f0a9"
   },
   "outputs": [],
   "source": [
    "# convert test loader into a list so that we may index it.\n",
    "train_loader_list = list(train_loader)\n",
    "\n",
    "plt.figure(figsize=(10,10))\n",
    "for i in range(25):\n",
    "    plt.subplot(5,5,i+1)\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n",
    "    plt.grid(False)\n",
    "    img = inverse_normalize(train_loader_list[0][0][i]).permute(0,2,1).contiguous().permute(2,1,0)\n",
    "    plt.imshow(img, cmap=plt.cm.binary)\n",
    "    plt.xlabel(class_names[train_loader_list[0][1][i]])\n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Qnjh12UbNFpV"
   },
   "source": [
    "## Networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "RGbLY6X-NH4O",
    "outputId": "018313f3-afa4-43ad-c72d-1f51fe5aedb6"
   },
   "outputs": [],
   "source": [
    "def initialize_params(module, mean, stdDev):\n",
    "    \"\"\"Helper function that initializes the biases and weights of a module given mean and standard deviations\n",
    "    given a module that is either a 2D Convolution or a 2D Transpose Convolution  to \n",
    "    \"\"\"\n",
    "    if isinstance(module, nn.ConvTranspose2d) or isinstance(module, nn.Conv2d):\n",
    "        module.weight.data.normal_(mean, stdDev)\n",
    "\n",
    "class Generator(nn.Module):\n",
    "    \"\"\"A deep convolutional Generator network, \n",
    "    designed to generate 32x32 colour images from a random 100D noise vector.\n",
    "    Based on DCGAN but with one less layer, \n",
    "    mostly due to the fact that the images here are half the size of those in the original paper.\n",
    "    \"\"\"\n",
    "    def __init__(self, b=128):\n",
    "        # inheriting initialization parameters from nn.Module\n",
    "        super(Generator, self).__init__()\n",
    "        # transpose convolve input into 512x4x4 tensor\n",
    "        self.tconv4 = nn.ConvTranspose2d(100, b*4, kernel_size = 4, stride = 1, padding = 0, bias = False)\n",
    "        self.bnorm4 = nn.BatchNorm2d(b*4)\n",
    "        # transpose convolve into 256x8x8 tensor\n",
    "        self.tconv8 = nn.ConvTranspose2d(b*4, b*2, kernel_size = 4, stride = 2, padding = 1, bias = False)\n",
    "        self.bnorm8 = nn.BatchNorm2d(b*2)\n",
    "        # transpose convolve into 128x16x16 tensor\n",
    "        self.tconv16 = nn.ConvTranspose2d(b*2, b, kernel_size = 4, stride = 2, padding = 1, bias = False)\n",
    "        self.bnorm16 = nn.BatchNorm2d(b)\n",
    "        # transpose convolve into 3x32x32 tensor (an image)\n",
    "        self.tconv32 = nn.ConvTranspose2d(b, 3, kernel_size = 4, stride = 2, padding = 1)\n",
    "        self.tanh = nn.Tanh()\n",
    "        \n",
    "        # activation function \n",
    "        self.relu = nn.ReLU()\n",
    "    \n",
    "    def initialize(self, mean, stdDev):\n",
    "        \"\"\"Instance method for initializing the weights to specific values\"\"\"\n",
    "        for module in self._modules:\n",
    "            initialize_params(self._modules[module], mean, stdDev)\n",
    "    \n",
    "    def forward(self, z):\n",
    "        \"\"\"Defines the order of operations to follow when the generator is called\"\"\"\n",
    "        z = self.relu(self.bnorm4(self.tconv4(z)))\n",
    "        z = self.relu(self.bnorm8(self.tconv8(z)))\n",
    "        z = self.relu(self.bnorm16(self.tconv16(z)))\n",
    "        z = self.tanh(self.tconv32(z))\n",
    "        return z\n",
    "\n",
    "class Discriminator(nn.Module):\n",
    "    \"\"\"A deep convolutional Discriminator network, \n",
    "    designed determine whether an input image is real (output 1) or fake (output 0).\n",
    "    Based on DCGAN but with one less layer, \n",
    "    mostly due to the fact that the images here are half the size of those in the original paper.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, b=128):\n",
    "        # inheriting initialization parameters from nn.Module\n",
    "        super(Discriminator, self).__init__()\n",
    "        # convolve image into 128x16x16 tensor\n",
    "        self.conv16 = nn.Conv2d(3, b, kernel_size = 4, stride = 2, padding = 1)\n",
    "        # convolve into 256x8x8 tensor\n",
    "        self.conv8 = nn.Conv2d(b, b*2, kernel_size = 4, stride = 2, padding = 1, bias = False)\n",
    "        self.bnorm8 = nn.BatchNorm2d(b*2)\n",
    "        # convolve into 512x4x4 tensor\n",
    "        self.conv4 = nn.Conv2d(b*2, b*4, kernel_size = 4, stride =2, padding = 1, bias = False)\n",
    "        self.bnorm4 = nn.BatchNorm2d(b*4)\n",
    "        # compute sigmoid\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "        # activation function\n",
    "        self.leaky_relu = nn.LeakyReLU(negative_slope=0.2)\n",
    "        \n",
    "    def initialize(self, mean, stdDev):\n",
    "        \"\"\"Instance method for initializing the weights to specific values\"\"\"\n",
    "        for module in self._modules:\n",
    "            initialize_params(self._modules[module], mean, stdDev)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \"\"\"Defines the order of operations to follow when the discriminator is called\"\"\"\n",
    "        x = self.leaky_relu(self.conv16(x))\n",
    "        x = self.leaky_relu(self.bnorm8(self.conv8(x)))\n",
    "        x = self.leaky_relu(self.bnorm4(self.conv4(x)))\n",
    "        x = self.sigmoid(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initiliaze the neural network instances, setting weights appropriately\n",
    "G = Generator().to(device)\n",
    "G.initialize(mean = 0, stdDev = 0.02)\n",
    "D = Discriminator().to(device)\n",
    "D.initialize(mean = 0, stdDev = 0.02)\n",
    "\n",
    "print(f'> Number of Generator parameters {len(torch.nn.utils.parameters_to_vector(G.parameters()))}')\n",
    "print(f'> Number of Discriminator parameters {len(torch.nn.utils.parameters_to_vector(D.parameters()))}')\n",
    "\n",
    "# define optimizer parameters\n",
    "learning_rate = 0.0002\n",
    "betas = (0.5, 0.999) # momentum\n",
    "\n",
    "# initialise the optimisers\n",
    "opt_G = torch.optim.Adam(G.parameters(), lr = learning_rate, betas = betas)\n",
    "opt_D = torch.optim.Adam(D.parameters(), lr = learning_rate, betas = betas )\n",
    "\n",
    "# set the epoch settings\n",
    "epoch = 0\n",
    "max_epochs = 100\n",
    "\n",
    "# initialize livelossplot instance\n",
    "liveplot = PlotLosses()\n",
    "\n",
    "# define loss function\n",
    "BCELoss = nn.BCELoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "N1UBl0PJjY-f"
   },
   "source": [
    "### Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 357
    },
    "colab_type": "code",
    "id": "kb5909Y8D_zx",
    "outputId": "f24a34b5-d39d-420b-cfd6-e0572b9d2fa5"
   },
   "outputs": [],
   "source": [
    "while epoch <= max_epochs:\n",
    "    g_losses = []\n",
    "    d_losses = []\n",
    "    \n",
    "    for (x, _) in train_loader:\n",
    "        # passing the tensor to to the current device \n",
    "        if str(device) == \"cuda\":\n",
    "            x = Variable(x.cuda())\n",
    "        else:\n",
    "            x.to(device) \n",
    "        \n",
    "        # define what ideal real and fake batches should output\n",
    "        real_out = torch.ones(1)[0].to(device)\n",
    "        fake_out = torch.zeros(1)[0].to(device)\n",
    "        \n",
    "        ############   Train Generator   #############\n",
    "        # set the gradient to zero\n",
    "        opt_G.zero_grad()\n",
    "        \n",
    "        # generate batch_size x 100 x 1 x 1 sample of random noise\n",
    "        random_noise = torch.rand(x.size(0), 100, 1, 1).to(device)\n",
    "        # utilize the random noise to generate images\n",
    "        gen_images = G(random_noise)\n",
    "        # pass the generated images to the discriminator, determines whether each image is real or false\n",
    "        dis_gen_outcome = D(gen_images).mean()\n",
    "        \n",
    "        # calculate the 'heauristic - non saturating game' binary cross entropy for the generator\n",
    "        gen_batch_loss = BCELoss(dis_gen_outcome, real_out)\n",
    "        # backpropagate loss\n",
    "        gen_batch_loss.backward()\n",
    "        # perform gradient decent\n",
    "        opt_G.step()\n",
    "        \n",
    "        ############ Train Discriminator #############\n",
    "        # set the gradient to zero\n",
    "        opt_D.zero_grad()\n",
    "        \n",
    "        # pass the real images to the discriminator, determines whether each image is real or false\n",
    "        dis_real_outcome = D(x).mean()\n",
    "        # calculate real discrimination loss\n",
    "        dis_batch_real_loss = BCELoss(dis_real_outcome, real_out)\n",
    "        \n",
    "        # generate batch_size x 100 x 1 x 1 sample of random noise\n",
    "        random_noise = torch.rand(x.size(0), 100, 1, 1).to(device)\n",
    "        # utilize the random noise to generate images\n",
    "        gen_images = G(random_noise)\n",
    "        # pass the generated images to the discriminator, determines whether each image is real or false\n",
    "        dis_gen_outcome = D(gen_images).mean()\n",
    "        # calculate fake discrimination loss\n",
    "        dis_batch_fake_loss = BCELoss(dis_gen_outcome, fake_out)\n",
    "        \n",
    "        # calculate the total discriminator loss\n",
    "        dis_batch_loss = (dis_batch_real_loss + dis_batch_fake_loss)/2.0\n",
    "        # backpropagate loss\n",
    "        dis_batch_loss.backward()\n",
    "        # perform gradient descent\n",
    "        opt_D.step()\n",
    "        ##############################################\n",
    "        \n",
    "        # record losses for this batch\n",
    "        g_losses.append(gen_batch_loss.cpu().data)\n",
    "        d_losses.append(dis_batch_loss.cpu().data)\n",
    "    \n",
    "    # plot results\n",
    "    plt.grid(False)\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n",
    "    plt.title(\"32 Images Generated by the Generator\")\n",
    "    plt.imshow(torchvision.utils.make_grid((gen_images[torch.randperm(128)])[:32], nrow = 8, normalize=True).cpu().data.permute(0,2,1).contiguous().permute(2,1,0), cmap=plt.cm.binary)\n",
    "  \n",
    "    liveplot.update({\n",
    "        'generator loss': np.mean(g_losses),\n",
    "        'discriminator loss': np.mean(d_losses)\n",
    "    })\n",
    "    liveplot.draw()\n",
    "    sleep(1.)\n",
    "\n",
    "    epoch = epoch+1\n",
    "\n",
    "print(\"Training completed. Saving model state dictionary\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# saving the model weights\n",
    "model_dictionary = G.state_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MTFqiEHzMOVw"
   },
   "source": [
    "## Results\n",
    "To generate the Pegasus, we try two approaches: An interpolation between a horse and a bird, and an interpolation between a horse and plane (a cyborg pegasus). \n",
    "\n",
    "To interpolate, we need vectors. We recover vectors of exemplar horse, plane and bird images using [Stochastic Clipping](https://arxiv.org/pdf/1702.04782.pdf)\n",
    "\n",
    "The interpolation is then simply vector arithmetic. Currently we simply interpolate halfway between the two recovered vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# defining the path, uncomment first line and comment second line if using google COLAB\n",
    "# path = \"/content/gdrive/My Drive/Education/Undergraduate/Year_3/Computer_Science/SSA/Machine_Learning/Coursework/ML_Classifier-Pegasus-Generator/Exemplary_Tensors/\"\n",
    "path = \"Exemplary_Tensors/\"\n",
    "\n",
    "# get exemplar tensors\n",
    "bird1 =  torch.load(path+\"bird1.pt\")\n",
    "bird2 =  torch.load(path+\"bird2.pt\")\n",
    "horse =  torch.load(path+\"horse.pt\")\n",
    "plane =  torch.load(path+\"plane.pt\")\n",
    "\n",
    "exemplar_tensors_list = [bird1, bird2, horse, plane]\n",
    "\n",
    "# show exemplar tensors\n",
    "plt.figure(figsize=(8,8))\n",
    "for i in range(4):\n",
    "    plt.subplot(2,2,i+1)\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n",
    "    plt.grid(False)\n",
    "    img =  inverse_normalize(exemplar_tensors_list[i]).permute(0,2,1).contiguous().permute(2,1,0)\n",
    "    plt.imshow(img, cmap=plt.cm.binary)\n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def recover_vector(target_image, generator):\n",
    "    \"\"\"Recovers an approximate latent vector z_approx corresponding to a given an original image G(z)\n",
    "    This is following https://arxiv.org/pdf/1702.04782.pdf\n",
    "    \"\"\"\n",
    "    # initialize random noise tensor (we will iteratively improve its components)\n",
    "    z_approx = torch.rand(128, 100, 1, 1).to(device)\n",
    "    \n",
    "    if str(device) == \"cuda\":\n",
    "        target_image = Variable(target_image.cuda())\n",
    "    else:\n",
    "        target_image.to(device) \n",
    "     \n",
    "    # instantiate loss function (L2 norm)\n",
    "    mse_loss = nn.MSELoss()\n",
    "    mse_loss.to(device)\n",
    "  \n",
    "    # instantiate optimizer \n",
    "    opt = torch.optim.Adam([z_approx], lr=0.01, betas=(0.05, 0.999))\n",
    "    \n",
    "    # minimize the difference\n",
    "    print(\"recovering vector\")\n",
    "    for i in range(100):\n",
    "        # set the gradient to zero\n",
    "        opt.zero_grad()\n",
    "        # generate the image from the noise\n",
    "        approx_image = generator(z_approx)\n",
    "        approx_image.to(device)\n",
    "        # calculate loss\n",
    "        loss = mse_loss(torch.mean(approx_image, dim=0), target_image)\n",
    "        # backpropagate\n",
    "        loss.backward()\n",
    "        # perform gradient descent\n",
    "        opt.step()\n",
    "        \n",
    "        # perform stochastic clipping\n",
    "        z_approx.data[z_approx.data > 1] = np.random.uniform(-1, 1)\n",
    "        z_approx.data[z_approx.data < -1] = np.random.uniform(-1, 1)\n",
    "        \n",
    "    return z_approx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# defining trained generator\n",
    "trained_G = Generator(b=128).to(device)\n",
    "# load the weights into this generator\n",
    "trained_G.load_state_dict(model_dictionary)\n",
    "\n",
    "# get the latent vectors for the birds\n",
    "bird_vector_1 = recover_vector(bird1, trained_G)\n",
    "bird_vector_2 = recover_vector(bird1, trained_G)\n",
    "# get the latent vector for a horse\n",
    "horse_vector = recover_vector(horse, trained_G)\n",
    "# get the latent vector for the plane\n",
    "plane_vector = recover_vector(plane, trained_G)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate bird-horse pegasus\n",
    "bird_horse_diff = bird_vector_2 - horse_vector \n",
    "bh_pegasus_vector = horse_vector + bird_horse_diff/2\n",
    "bh_pegasus_tensor = torch.mean(trained_G(bh_pegasus_vector), dim=0)\n",
    "bh_pegasus_image = inverse_normalize(bh_pegasus_tensor.cpu()).permute(0,2,1).contiguous().permute(2,1,0)\n",
    "\n",
    "# generate plane-horse pegasus\n",
    "plane_horse_diff = plane_vector - horse_vector \n",
    "ph_pegasus_vector = horse_vector + plane_horse_diff/2\n",
    "ph_pegasus_tensor = torch.mean(trained_G(ph_pegasus_vector), dim=0)\n",
    "ph_pegasus_image = inverse_normalize(ph_pegasus_tensor.cpu()).permute(0,2,1).contiguous().permute(2,1,0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8,8))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.title(\"Bird-Horse Pegasus\")\n",
    "plt.xticks([])\n",
    "plt.yticks([])\n",
    "plt.grid(False)\n",
    "plt.imshow(bh_pegasus_image.detach().numpy())\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(\"Plane-Horse Pegasus\")\n",
    "plt.xticks([])\n",
    "plt.yticks([])\n",
    "plt.grid(False)\n",
    "plt.imshow(ph_pegasus_image.detach().numpy())\n",
    "plt.tight_layout()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "coursework-pegasus.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}