Day 4 - Autoencoders.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Autoencoders and Variational Autoencoders (VAEs)\n",
    "\n",
    "## Motivation\n",
    "\n",
    "- most real world data is unlabelled, and we would like to be able to learn from it and use it\n",
    "- a lot of data in inputs is not useful. How can we find the importance within it?\n",
    "\n",
    "When we see things, we don't consciously percieve the wavelength of light hitting each of our individual cells. This is the very high dimensional, raw input data that your eyes recieve, but this information is too complex for us to use for planning and decision making etc. Instead, what we do is perform abstraction from it - we build up more meaningful, general and vague descriptions of the environment.\n",
    "We percieve and would describe the scene using a few high level *features* such as: ' I see a face', 'It is smiling', 'It is nearby' rather than the millions of raw wavelengths. We somehow manage to map the high dimensional, low level wavelength data into much lower dimensional, high level concepts. The variables that we choose to represent the data in this lower dimensional space are called **latent variables**. \n",
    "\n",
    "## What is an autoencoder?\n",
    "\n",
    "An autoencoder is a unsupervised machine learning model that can learn a latent representation of input data, that is a lower dimensional higher level \n",
    "An autoencoder learns to do so by passing the input through a neural network which attempts to output the same thing (or slightly different) as it is given as an input, whilst passing the data through a low dimensional bottleneck in between. We can determine a loss by taking the BCE loss between the input and its reconstruction at the output. This means that training requires no labels!\n",
    "\n",
    "<img src=\"autoencoder.png\">\n",
    "\n",
    "A lower dimensional representation is much more easily interpretable than the higher dimensional input, as there is just less information to comprehend.\n",
    "Having a trained autoencoder means that we have found a way to represent our data in a lower dimension with minimal loss.\n",
    "Our inputs can be reconstructed \n",
    "\n",
    "We will train out autoencoder to come up with valuable ways of representing the data in interpretable, low-dimensional space, by attempting to minimise the binary cross entropy loss between the input image and its reconstruction.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch    # computational library saves us time and effort in building models\n",
    "from torch.autograd import Variable     # for computational graphs\n",
    "import matplotlib.pyplot as plt     # for plotting\n",
    "from torchvision import datasets, transforms    # to get the dataset and then transform it into a tensor\n",
    "import torch.nn.functional as F # functional stuff like our activation function\n",
    "import numpy as np\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from time import sleep\n",
    "\n",
    "##%matplotlib inline\n",
    "#%matplotlib notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 1    # for how many runs through of the dataset?\n",
    "lr = 0.001      # proportionality constant controlling parameter update size\n",
    "batch_size = 128 # how large are the training batches?\n",
    "latent_dim = 2  # how many dimensions does our latent variable have? we can plot it in 3\n",
    "\n",
    "train_data = datasets.FashionMNIST(root='fashiondata/',     # where to save/look for it\n",
    "                         train=True,    # this is for training\n",
    "                         transform=transforms.ToTensor(),   # transform it into a tensor of data\n",
    "                         download=True) # yes, download it\n",
    "\n",
    "test_data = datasets.FashionMNIST(root='fashiondata/',\n",
    "                           train=False,\n",
    "                           transform=transforms.ToTensor(),\n",
    "                           download=True)\n",
    "\n",
    "# make a dataloader to generate us samples for training\n",
    "training_samples = torch.utils.data.DataLoader(dataset=train_data,\n",
    "                                               batch_size=batch_size,   # train on small batches of inputs\n",
    "                                               shuffle=True)    # make sure to shuffle the data to avoid overfitting\n",
    "\n",
    "testing_samples = torch.utils.data.DataLoader(dataset=test_data,\n",
    "                                           batch_size=batch_size,\n",
    "                                           shuffle=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class autoencoder(torch.nn.Module): # create a class for our autoencoder\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        # encoder\n",
    "        self.e1 = torch.nn.Linear(784, 256)\n",
    "        self.e2 = torch.nn.Linear(256, 64)\n",
    "        self.e3 = torch.nn.Linear(64, latent_dim)\n",
    "\n",
    "        # decoder\n",
    "        self.d1 = torch.nn.Linear(latent_dim, 64)\n",
    "        self.d2 = torch.nn.Linear(64, 256)\n",
    "        self.d3 = torch.nn.Linear(256, 784)\n",
    "\n",
    "    def encode(self, x): # compress our input into the latent space\n",
    "        x = x.view(-1, 784)             # unroll\n",
    "        z = F.relu(self.e1(x))\n",
    "        z = F.relu(self.e2(z))\n",
    "        z = self.e3(z)\n",
    "        #print(z)\n",
    "        return z\n",
    "\n",
    "    def decode(self, z):\n",
    "        x_pred = F.relu(self.d1(z))\n",
    "        x_pred = F.relu(self.d2(x_pred))\n",
    "        x_pred = F.sigmoid(self.d3(x_pred))\n",
    "        return x_pred\n",
    "\n",
    "    def forward(self, x):   # decode\n",
    "        z = self.encode(x)\n",
    "        x_pred = self.decode(z)\n",
    "        return x_pred, z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The loss function we are going to use is the binary cross entropy between the pixels of the original image and its reconstruction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss(x_hat, x):\n",
    "    reconstruction_loss = F.binary_cross_entropy(x_hat, x.view(-1, 784))\n",
    "    return reconstruction_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "myautoencoder = autoencoder()\n",
    "optimiser = torch.optim.Adam(myautoencoder.parameters(), lr=lr)#, weight_decay=0.0001)\n",
    "\n",
    "def train():\n",
    "    myautoencoder.train()   # put in training mode\n",
    "    costs = []\n",
    "    for i in range(epochs):\n",
    "        for batch_index, (x, y) in enumerate(training_samples):\n",
    "\n",
    "            x = Variable(x)\n",
    "            #print('x shape', x.shape)\n",
    "            x_pred, z = myautoencoder(x)  # (calls myvae.forward) generate an output\n",
    "\n",
    "            cost = loss(x_pred, x)\n",
    "            costs.append(cost.data)\n",
    "            cost.backward()\n",
    "            optimiser.step()\n",
    "            optimiser.zero_grad()\n",
    "            print('batch', batch_index, 'cost', cost.data[0])\n",
    "            \n",
    "           \n",
    "        \n",
    "            z = np.array(z.data)\n",
    "            #print(z.shape)\n",
    "            colordict = {0:'blue', 1:'orange', 2:'green', 3:'red', 4:'purple', 5:'brown', 6:'pink', 7:'gray', 8:'olive', 9:'cyan'}\n",
    "            colorlist = [colordict[i] for i in y]\n",
    "            \n",
    "            plt.figure()\n",
    "            plt.scatter(z[:, 0], z[:, 1], c=colorlist, s=5)\n",
    "            plt.show()\n",
    "            \n",
    "            plt.figure()\n",
    "            plt.plot(costs, 'b')\n",
    "            plt.show()\n",
    "\n",
    "            if batch_index == 200:\n",
    "                break\n",
    "                \n",
    "            if batch_index % 20:\n",
    "                x = x.view(-1, 28, 28)\n",
    "                x_pred = x_pred.view(-1, 28, 28)\n",
    "\n",
    "                \n",
    "                \n",
    "                plt.imshow(x.data[0])\n",
    "                plt.show()\n",
    "                plt.imshow(x_pred.data[0])\n",
    "                plt.show()\n",
    "                \n",
    "train()\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "plt.ioff()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "### Use cases of autoencoders\n",
    "#### Denoising\n",
    "Autoencoders can be used to remove noise from images by training them to reconstruct images from those which have had noise added to them. In a similar way, they can also be used to remove watermarks or add colour to images.\n",
    "\n",
    "<img src='denoising autoencoder.png'>\n",
    "\n",
    "#### Learning latent representations\n",
    "\n",
    "As we have seen, autoencoders can find a way to project input data down into a lower dimension whilst keeping it useful. They can extract higher level concepts from low dimensional data - essentially abstracting from the data.\n",
    "\n",
    "**Autoencoders allow us to learn the higher level features within the raw data without using labels!**\n",
    "\n",
    "### Problems with autoencoders\n",
    "\n",
    "- the reconstruction of the input will be lossy, as information must be thrown away to compress the input down to a lower dimension. Many other techniques, like JPEG, far outperform autoencoders in terms of image comp\n",
    "- the weights learned by the model are specific to the data it has been trained on. There is nothing preventing  a lot of latent space never being sampled for training. As such, the latent space may be sparse and can not be sampled from to produce an outpput. This prevents autoencoders being used as generative models.\n",
    "- the diverging part of the model can not robustly be used to generate data, because the latent space may be filled with gaps where that model has not learnt to reconstruct any data from. \n",
    "- the latent representations for similar inputs/inputs may not be nearby in the latent space.\n",
    "\n",
    "## Variational Autoencoders\n",
    "\n",
    "Variational autencoders are an upgrade on autoencoders. The VAE model differs from that of a standard autoencoder in two ways.\n",
    "\n",
    "1. The latent representation  of a given input,which we reconstruct from, is stochastic. The encoder constists of two networks. The output of these two networks are two vectors which represent means and log(variances) respectively, in each dimension, for each input which we will use to sample to latent vector from. Now the latent vector is non-deterministic and leads to the model sampling from locations all around where the latent vector is centered - so the latent space becomes filled around the means. \n",
    "\n",
    "It's corresponding latent representation is then produced by sampling from a Gaussian characterised by these parameters. We are not learning *these* parameters, but the weights in the encoder network leading to them, which must make them useful for reconstructing the input later.\n",
    "\n",
    "2. The loss function includes a KL-divergence term. KL-divergence is used to make distributions more similar; where functions have greatly different outputs, the KL-divergence will be large, and where they are the same, the KL-divergence will be zero. In our case we take the KL-divergence between our generated probability distributions, parameterised by our predicted means and variances, and a unit Gaussian (our prior, which we want to make our data distributed like. I.e. This part of the loss function forces distributions of the latent variables for each class to be centred around the origin and have unit variance. \n",
    "\n",
    "These two things allow us to use the model generatively, because the model has experience with what the reconstruction should be for a much greater range of latent inputs.\n",
    "\n",
    "For a unit Gaussian prior, the graph of how the KL-divergence varies with variance (x) and mean (y) is show below:\n",
    "<img src='KL.png'>\n",
    "\n",
    "So our loss function becomes:\n",
    "\n",
    "$ J = \\frac{1}{2}\\sum(h - y)^2 -\\frac{1}{2}\\sum (1 + log(\\sigma^2) - \\sigma^2 - \\mu^2)$\n",
    "\n",
    "The first term is the reconstruction loss, and the second term is the KL-divergence. They balance reconstruction and \n",
    "\n",
    "They ensure that the latent space is filled, and that similar outputs occupy nearby regions in the latent space. They do this by minimising the difference between how the latent variables are distributed compared wth some prior (initially expected) distribution, $p(z)$. We usually take our prior distribution over $z$ to be a unit Gaussian - a reasonable assumption, as lots of things in real life are distributed normally.\n",
    "So each latent variable now is forced to be distributed normally around zero, with unit variance. This \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch    # computational library saves us time and effort in building models\n",
    "from torch.autograd import Variable     # for computational graphs\n",
    "import matplotlib.pyplot as plt     # for plotting\n",
    "from torchvision import datasets, transforms    # to get the dataset and then transform it into a tensor\n",
    "import torch.nn.functional as F # functional stuff like our activation function\n",
    "import numpy as np\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from time import sleep\n",
    "\n",
    "epochs = 1    # for how many runs through of the dataset?\n",
    "lr = 0.02      # proportionality constant controlling parameter update size\n",
    "batch_size = 128 # how large are the training batches?\n",
    "latent_dim = 40  # how many dimensions does our latent variable have? we can plot it in 3\n",
    "beta = 1\n",
    "\n",
    "train_data = datasets.FashionMNIST(root='fashiondata/',     # where to save/look for it\n",
    "                         train=True,    # this is for training\n",
    "                         transform=transforms.ToTensor(),   # transform it into a tensor of data\n",
    "                         download=True) # yes, download it\n",
    "\n",
    "test_data = datasets.FashionMNIST(root='fashiondata/',\n",
    "                           train=False,\n",
    "                           transform=transforms.ToTensor(),\n",
    "                           download=True)\n",
    "\n",
    "# make a dataloader to generate us samples for training\n",
    "training_samples = torch.utils.data.DataLoader(dataset=train_data,\n",
    "                                               batch_size=batch_size,   # train on small batches of inputs\n",
    "                                               shuffle=True)    # make sure to shuffle the data to avoid overfitting\n",
    "\n",
    "testing_samples = torch.utils.data.DataLoader(dataset=test_data,\n",
    "                                           batch_size=batch_size,\n",
    "                                           shuffle=True)\n",
    "\n",
    "class VAE(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        # encoder to means and standard deviations\n",
    "        self.to_mu1 = torch.nn.Linear(784, 256)\n",
    "        self.to_mu2 = torch.nn.Linear(256, 64)\n",
    "        self.to_mu3 = torch.nn.Linear(64, latent_dim)\n",
    "\n",
    "        self.to_logvar1 = torch.nn.Linear(784, 256)\n",
    "        self.to_logvar2 = torch.nn.Linear(256, 64)\n",
    "        self.to_logvar3 = torch.nn.Linear(64, latent_dim)\n",
    "\n",
    "        # decoder\n",
    "        self.d1 = torch.nn.Linear(latent_dim, 64)\n",
    "        self.d2 = torch.nn.Linear(64, 256)\n",
    "        self.d3 = torch.nn.Linear(256, 784)\n",
    "\n",
    "    def encode(self, x):\n",
    "        x = x.view(-1, 784)\n",
    "\n",
    "        mu = F.relu(self.to_mu1(x))\n",
    "        mu = F.relu(self.to_mu2(mu))\n",
    "        mu = self.to_mu3(mu)\n",
    "\n",
    "        logvar = F.relu(self.to_logvar1(x))\n",
    "        logvar = F.relu(self.to_logvar2(logvar))\n",
    "        logvar = self.to_logvar3(logvar)\n",
    "\n",
    "        return mu, logvar\n",
    "\n",
    "    def reparameterize(self, mu, logvar):\n",
    "        epsilon = Variable(torch.Tensor(np.random.randn(batch_size, latent_dim)))\n",
    "        z = mu + epsilon * (0.5*logvar).exp()\n",
    "        return z\n",
    "\n",
    "    def decode(self, z):\n",
    "        x_pred = F.relu(self.d1(z))\n",
    "        x_pred = F.relu(self.d2(x_pred))\n",
    "        x_pred = F.sigmoid(self.d3(x_pred))\n",
    "        #print(x_pred.shape)\n",
    "        return x_pred\n",
    "\n",
    "    def forward(self, x):\n",
    "        mu, logvar = self.encode(x) # the output for the sd does not make sense to be negative, so predict log(var)\n",
    "        z = self.reparameterize(mu, logvar)\n",
    "        x_pred = self.decode(z)\n",
    "        return x_pred, z, mu, logvar\n",
    "\n",
    "def VAEloss(x_hat, x, mu, logvar):\n",
    "    reconstruction_loss = F.binary_cross_entropy(x_hat, x.view(-1, 784), size_average=False)\n",
    "    #print(reconstruction_loss.data[0])\n",
    "    KL_divergence = - 0.5 * torch.sum(1 + logvar - mu**2 - logvar.exp())\n",
    "    #print(KL_divergence.data[0])\n",
    "    return reconstruction_loss + beta * KL_divergence\n",
    "\n",
    "fig1 = plt.figure(figsize=(10, 20))\n",
    "ax1 = fig1.add_subplot(121)\n",
    "ax1.set_title('Costs vs epoch')\n",
    "#ax1.set_ylim(0, 500)\n",
    "ax2 = fig1.add_subplot(122)\n",
    "ax2.set_title('Latent representation')\n",
    "\n",
    "fig2 = plt.figure(figsize=(10, 20))\n",
    "ax3 = fig2.add_subplot(121)\n",
    "ax3.set_title('Input')\n",
    "ax4 = fig2.add_subplot(122)\n",
    "ax4.set_title('Reconstruction')\n",
    "\n",
    "plt.ion()\n",
    "plt.show()\n",
    "\n",
    "myVAE = VAE()\n",
    "optimiser = torch.optim.Adam(myVAE.parameters(), lr=lr)\n",
    "\n",
    "def trainVAE():\n",
    "    myVAE.train()   # put in training mode\n",
    "    costs = []\n",
    "    for epoch in range(epochs):\n",
    "        for batch_index, (x, y) in enumerate(training_samples):\n",
    "\n",
    "            x = Variable(x)\n",
    "            #print('x shape', x.shape)\n",
    "            #optimiser.zero_grad()\n",
    "            x_pred, z, mu, logvar = myVAE(x)  # (calls myvae.forward) generate an output\n",
    "\n",
    "            cost = VAEloss(x_pred, x, mu, logvar)\n",
    "            costs.append(cost.data)\n",
    "            cost.backward()\n",
    "            optimiser.step()\n",
    "            optimiser.zero_grad()\n",
    "            print('Epoch', epoch, 'batch', batch_index, 'cost', cost.data[0])\n",
    "\n",
    "            z = np.array(z.data) # for plotting\n",
    "            colordict = {0:'blue', 1:'orange', 2:'green', 3:'red', 4:'purple', 5:'brown', 6:'pink', 7:'gray', 8:'olive', 9:'cyan'}\n",
    "            colorlist = [colordict[i] for i in y]\n",
    "            ax2.scatter(z[:, 0], z[:, 1], c=colorlist, s=5)\n",
    "            ax1.plot(costs, 'b')\n",
    "            fig1.canvas.draw()\n",
    "\n",
    "            if batch_index % 20:\n",
    "                x = x.view(-1, 28, 28)\n",
    "                x_pred = x_pred.view(-1, 28, 28)\n",
    "\n",
    "                ax3.imshow(x.data[0])\n",
    "                ax4.imshow(x_pred.data[0])\n",
    "\n",
    "                fig2.canvas.draw()\n",
    "\n",
    "            if batch_index == 30:\n",
    "                ax2.clear()\n",
    "\n",
    "            if batch_index == 200:\n",
    "                #return ax2\n",
    "                break\n",
    "\n",
    "\n",
    "trainVAE()\n",
    "torch.save(myVAE, 'trainedVAE.pt')\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "#plt.ioff()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Problems with VAEs\n",
    "\n",
    "The latent factors found are often entangled. The latent vectors do not represent variation of only one high level feature. Beta VAEs try to confront this problem and put more of a penalty on the latent variables that are found being \n",
    "\n",
    "\n",
    "## Beta Variational Autencoders\n",
    "\n",
    "So for VAEs, we have two terms in the loss function: the reconstruction loss which makes sure that we get back a similar image to our input. And the KL-divergence, which ensures that the latent space which a realistic output can be generated from is continuous, and that similar examples are nearby in the latent space. \n",
    "Beta-VAEs use another hyperparameter (Beta) to scale the KL-divergence part of the loss"
   ]
  }
 ],
 "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
}