code/vae_notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# How to implement a Variational Autoencoder (VAE)\n",
    "\n",
    "A variational autoencoder observes data, infers a latent code for it and tries to reconstruct the data from that latent code. In contrast to regular autoencoders, the code of the VAE is **random**. That means that when presented with the same input, the VAE will produce a slightly different code each time. This makes its decoding process more robust, since it has to deal with noisy code.\n",
    "\n",
    "Another way of looking at a VAE is as a training procedure for a probablistic model. The model is \n",
    "$$p(x) = \\int p(z)p(x|z) dz$$\n",
    "where $z$ is the latent code and $x$ is the data. During training we need to infer a posterior over $z$. In the case of a VAE this is done by neural network.\n",
    "\n",
    "Assuming that the theory of VAEs has already been presented, we now dive straight into implementing them. If you need more background on VAEs, have a look at our [tutorial slides](https://github.com/philschulz/VITutorial/tree/master/modules) and the references therein."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Framework\n",
    "\n",
    "For the purpose of this tutorial we are going to use [mxnet](https://mxnet.incubator.apache.org) which is a scalable deep learning library that has interfaces for several languages, including python. We are going to import and abbreviate it as \"mx\". We will use mxnet to define a computation graph. This is done using the [symbol library](https://mxnet.incubator.apache.org/api/python/symbol.html). When building the VAE, all the methods that you use should be prefixed with `mx.sym`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mxnet as mx\n",
    "import numpy as np\n",
    "import urllib.request\n",
    "import os, logging, sys\n",
    "from os.path import join, exists\n",
    "from abc import ABC\n",
    "from typing import List, Tuple, Callable, Optional, Iterable\n",
    "from matplotlib import cm, pyplot as plt\n",
    "from math import sqrt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we specify a couple of constants that will help us to load the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEFAULT_LEARNING_RATE = 0.0003\n",
    "\n",
    "TRAIN_SET = 'train'\n",
    "VALID_SET = 'valid'\n",
    "TEST_SET = 'test'\n",
    "data_names = [TRAIN_SET, VALID_SET, TEST_SET]\n",
    "test_set = [TEST_SET]\n",
    "data_dir = join(os.curdir, \"binary_mnist\")\n",
    "\n",
    "# change this to mx.gpu(0) if you want to run your code on gpu\n",
    "ctx = mx.cpu()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we set up basic logging facilities to print intermediate output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logging.basicConfig(level=logging.DEBUG, format=\"%(asctime)s [%(levelname)s]: %(message)s\", datefmt=\"%H:%M:%S\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Data\n",
    "\n",
    "Throughout the tutorial we will use the binarised MNIST data set consisting of images of handwritten digits (0-9). \n",
    "The binarisation was done by sampling each pixel from a Bernoulli distribution with the pixel's original intensity being the Bernoulli parameter (see [this paper](http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.141.1680&rep=rep1&type=pdf) for details).Each pixel has been mapped to either 0 or 1, meaning that pixels are either fully on or off. We use this data set because it allows us to use a rather simple product of Bernoullis as a likelihood. We download the data into a folder called \"binary_mnist\". This may take some time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not exists(data_dir):\n",
    "    os.mkdir(data_dir)\n",
    "for data_set in data_names:\n",
    "    file_name = \"binary_mnist.{}\".format(data_set)\n",
    "    goal = join(data_dir, file_name)\n",
    "    if exists(goal):\n",
    "        logging.info(\"Data file {} exists\".format(file_name))\n",
    "    else:\n",
    "        logging.info(\"Downloading {}\".format(file_name))\n",
    "        link = \"http://www.cs.toronto.edu/~larocheh/public/datasets/binarized_mnist/binarized_mnist_{}.amat\".format(\n",
    "            data_set)\n",
    "        urllib.request.urlretrieve(link, goal)\n",
    "        logging.info(\"Finished\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alright, now we have the data on disk. We will load the training and test data later. Right now we just load the validation set to see what the data looks like. We'll grab 5 random numbers and visualise them (you can run this box several times to get 5 different digits each time). That the digits look a bit rough results from their pixels being binarised."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_name = join(data_dir, \"binary_mnist.{}\".format(VALID_SET))\n",
    "logging.info(\"Reading {} into memory\".format(file_name))\n",
    "valid_set = np.genfromtxt(file_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random_indeces = np.random.randint(valid_set.shape[0], size=5)\n",
    "random_pictures = valid_set[random_indeces, :]\n",
    "width = height = int(sqrt(random_pictures.shape[1]))\n",
    "plot, axes = plt.subplots(1,5,  sharex='col', sharey='row', figsize=(15,3))\n",
    "for i in range(5):\n",
    "    axes[i].imshow(np.reshape(random_pictures[i,:],(width,height)), cmap=cm.Greys)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After having acquainted ourselves with the data we can now proceed to build our VAE."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Diagonal Gaussian VAE\n",
    "\n",
    "The most basic VAE model is one where we assume that the latent variable is multiviariate Gaussian. We fix the prior to be standard normal. During inference, we use a multivariate Gaussian variational distribution with diagonal covariance matrix. This means that we are only modelling variance but not covariance (in fact, a k-dimensional Guassian with diagonal covariance has the same density as a product of k independent univariate Gaussians). Geometrically, this variational distribution can only account for axis-aligned elliptical densities. It is thus rather limited in its modelling capabilities. Still, because it uses a neural network under the hood, it is very expressive. \n",
    "\n",
    "In this tutorial, we will model the mist binarised digit data set. Each image is encoded as a 784-dimensional vector. We will model each of these vectors as a product of 784 Bernoullis (of course, there are better models but we want to keep it simple). Our likelihood is thus a product of independent Bernoullis. The resulting model is formally specified as \n",
    "\n",
    "\\begin{align}z \\sim \\mathcal{N}(0,I) && x_i|z \\sim Bernoulli(NN_{\\theta}(z))~~~ i \\in \\{1,2,\\ldots, 784\\} \\ .\\end{align}\n",
    "\n",
    "The variational approximation is given by $$q(z|x) = \\mathcal{N}(NN_{\\lambda}(x), NN_{\\lambda}(x)).$$\n",
    "\n",
    "Notice that both the Bernoulli likelihood and the Gaussian variational distribution use NNs to compute their parameters. The parameters of the NNs, however, are different ($\\theta$ and $\\lambda$, respectively)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementation\n",
    "\n",
    "We will spread our implementation across 3 classes. This design choice is motivated by the desire to make our models as modular as possible. This will later allow us to mix and match different likelihoods and variational distributions.\n",
    "\n",
    "* **Generator**: This class defines our likelihood. Given a latent value, it produces a data sample or assigns a density to an existing data point.\n",
    "* **InferenceNetwork**: This neural network computes the parameters of the variational approximation from a data point.\n",
    "* **VAE**: This is the variational autoencoder. It combines a Generator and an InferenceNetwork and trains them jointly. Once trained, it can generate random data points or try to reproduce data presented to it.\n",
    "\n",
    "Below we have specified these classes abstractly. Make sure you understand what each method is supposed to be doing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator(ABC):\n",
    "    \"\"\"\n",
    "    Generator network.\n",
    "\n",
    "    :param data_dims: Dimensionality of the generated data.\n",
    "    :param layer_sizes: Size of each hidden layer in the network.\n",
    "    :param act_type: The activation after each hidden layer.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, data_dims: int, layer_sizes: List[int], act_type: str) -> None:\n",
    "        self.data_dims = data_dims\n",
    "        self.layer_sizes = layer_sizes\n",
    "        self.act_type = act_type\n",
    "\n",
    "    def generate_sample(self, latent_state: mx.sym.Symbol) -> mx.sym.Symbol:\n",
    "        \"\"\"\n",
    "        Generate a data sample from a latent state.\n",
    "\n",
    "        :param latent_state: The latent input state.\n",
    "        :return: A data sample.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def train(self, latent_state: mx.sym.Symbol) -> mx.sym.Symbol:\n",
    "        \"\"\"\n",
    "        Train the generator from a given latent state.\n",
    "        \n",
    "        :param latent_state: The latent input state.\n",
    "        :return: The loss symbol used for training.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "        \n",
    "class InferenceNetwork(ABC):\n",
    "    \"\"\"\n",
    "    A network to infer distributions over latent states.\n",
    "\n",
    "    :param latent_variable_size: The dimensionality of the latent variable.\n",
    "    :param layer_sizes: Size of each hidden layer in the network.\n",
    "    :param act_type: The activation after each hidden layer.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, latent_variable_size: int, layer_sizes: List[int], act_type: str) -> None:\n",
    "        self.latent_var_size = latent_variable_size\n",
    "        self.layer_sizes = layer_sizes\n",
    "        self.act_type = act_type\n",
    "\n",
    "    def inference(self, data: mx.sym.Symbol) -> Tuple[mx.sym.Symbol, ...]:\n",
    "        \"\"\"\n",
    "        Infer the parameters of the distribution over latent values.\n",
    "\n",
    "        :param data: A data sample.\n",
    "        :return: The parameters of the distribution.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "        \n",
    "    def get_kl_loss(params_q: Iterable[mx.sym.Symbol], params_p: Iterable[mx.sym.Symbol]) -> mx.sym.Symbol:\n",
    "        \"\"\"\n",
    "        Create a loss symbol for the kl term associated with the inference net.\n",
    "        \n",
    "        :param params_q: The parameters of the variational distribution.\n",
    "        :param params_p: The parameters of the model prior.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "        \n",
    "        \n",
    "class VAE(ABC):\n",
    "    \"\"\"\n",
    "    A variational autoencoding model (Kingma and Welling, 2014).\n",
    "\n",
    "    :param generator: A generator network that specifies the likelihood of the model.\n",
    "    :param inference_net: An inference network that specifies the distribution over latent values.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, generator: Generator, inference_net: InferenceNetwork) -> None:\n",
    "        self.generator = generator\n",
    "        self.inference_net = inference_net\n",
    "\n",
    "    def train(self, data: mx.sym.Symbol, label: mx.sym.Symbol) -> mx.sym.Symbol:\n",
    "        \"\"\"\n",
    "        Train the generator and inference network jointly by optimising the ELBO.\n",
    "\n",
    "        :param data: The training data.\n",
    "        :param label: Copy of the training data.\n",
    "        :return: A grouped list of loss symbols.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def generate_reconstructions(self, data: mx.sym.Symbol, n: int) -> mx.sym.Symbol:\n",
    "        \"\"\"\n",
    "        Generate a number of reconstructions of input data points.\n",
    "\n",
    "        :param data: The input data.\n",
    "        :param n: Number of reconstructions per data point.\n",
    "        :return: The reconstructed data.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def phantasize(self, n: int) -> mx.sym.Symbol:\n",
    "        \"\"\"\n",
    "        Generate data by randomly sampling from the prior.\n",
    "\n",
    "        :param n: Number of sampled data points.\n",
    "        :return: Randomly generated data points.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 1\n",
    "Let us start by implementing the generator. This is pretty much a standard neural network. The main point of this exercise is to get you comfortable with mxnet. Complete all the TODOs below. Before starting, check mxnet's [fully connected layer](https://mxnet.incubator.apache.org/api/python/symbol/symbol.html#mxnet.symbol.FullyConnected) and [activation functions](https://mxnet.incubator.apache.org/api/python/symbol.html#mxnet.symbol.Activation)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ProductOfBernoullisGenerator(Generator):\n",
    "    \"\"\"\n",
    "    A generator that produces binary vectors whose entries are independent Bernoulli draws.\n",
    "\n",
    "    :param data_dims: Dimensionality of the generated data.\n",
    "    :param layer_sizes: Size of each hidden layer in the network.\n",
    "    :param act_type: The activation after each hidden layer.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, data_dims: int, layer_sizes=List[int], act_type=str) -> None:\n",
    "        super().__init__(data_dims, layer_sizes, act_type)\n",
    "        # TODO choose the correct output activation for a Bernoulli variable. This should just be a string.\n",
    "        self.output_act = None\n",
    "\n",
    "    def _preactivation(self, latent_state: mx.sym.Symbol) -> mx.sym.Symbol:\n",
    "        \"\"\"\n",
    "        Computes the pre-activation of the generator, i.e. the hidden state before the final output activation.\n",
    "\n",
    "        :param latent_state: The input latent state.\n",
    "        :return: The pre-activation before output activation.\n",
    "        \"\"\"\n",
    "        prev_out = latent_state\n",
    "        for i, size in enumerate(self.layer_sizes):\n",
    "            # TODO for each layer size in layer_sizes, introduce a layer that performs an affine transformation \n",
    "            # (implemented by mx.sym.FullyConnected) followed by a non-linearity that is specified by self.act_type.\n",
    "            # Call the resulting symbol act_i and uncomment the line below\n",
    "            #prev_out = act_i\n",
    "            \n",
    "            # TODO: remove when finished\n",
    "            pass\n",
    "\n",
    "        # The output layer that gives pre_activations for multiple Bernoulli softmax between 0 and 1\n",
    "        fc_out = mx.sym.FullyConnected(data=prev_out, num_hidden=self.data_dims, name=\"gen_fc_out\")\n",
    "\n",
    "        return fc_out\n",
    "    \n",
    "    def generate_sample(self, latent_state: mx.sym.Symbol, binarise: bool = False) -> Tuple[mx.sym.Symbol, mx.sym.Symbol]:\n",
    "        \"\"\"\n",
    "        Generates a data sample. The stochasticity in the sampling\n",
    "        process comes from the latent_state. Either returns the Bernoulli parameters or the maximally likely\n",
    "        outcome.\n",
    "\n",
    "        :param latent_state: The input latent state.\n",
    "        :param binarise: Return discrete values instead of Bernoulli parameters.\n",
    "        :return: A vector of Bernoulli draws and the latent state from which they were generated.\n",
    "        \"\"\"\n",
    "        act = mx.sym.Activation(data=self._preactivation(latent_state=latent_state), act_type=self.output_act,\n",
    "                                name=\"gen_act_out\")\n",
    "        out = act > 0.5 if binarise else act\n",
    "        return out, latent_state\n",
    "    \n",
    "    def train(self, latent_state=mx.sym.Symbol, label=mx.sym.Symbol) -> mx.sym.Symbol:\n",
    "        \"\"\"\n",
    "        Train the generator from a given latent state.\n",
    "\n",
    "        :param latent_state: The input latent state.\n",
    "        :param label: A binary vector (same as input for inference module).\n",
    "        :return: The loss symbol used for training.\n",
    "        \"\"\"\n",
    "        output = mx.sym.Activation(data=self._preactivation(latent_state=latent_state), act_type=self.output_act,\n",
    "                                   name=\"output_act\")\n",
    "        return mx.sym.MakeLoss(-mx.sym.sum(label * mx.sym.log(output) + (1-label) * mx.sym.log(1-output), axis=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now move on to the inference network. Recall that this network will return the parameters of a diagonal Gaussian. Thus, we need to return two vectors of the same size: a mean and a standard deviation vector. (Formally, the parameters of the Gaussian are the variances. However, from the derivation of the Gaussian reparametrisation we know that we\n",
    "need the standard deviations to generate a Gaussian random variable $z$ as transformation of a standard Gaussian variable $\\epsilon$.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 2.a\n",
    "To train our VAE we need to compute the KL of the prior distribution from the variational approximation. Assuming that\n",
    "the prior is standard normal the KL for the a diagonal Gaussian variational approximation is\n",
    "$$ 0.5 * \\left(\\sum_i - \\log(\\sigma_i^2) - 1 + \\sigma_i^2 + \\mu_i^2 \\right) \\ . $$\n",
    "We use this KL divergence in the get_kl_loss method below of the GaussianInferenceNetwork. (Losses in mxnet are defined using the [MakeLoss](https://mxnet.incubator.apache.org/api/python/symbol/symbol.html#mxnet.symbol.MakeLoss) symbol.) Please implement the function diagonal_gaussian_kl. Use the [sum](https://mxnet.incubator.apache.org/api/python/symbol/symbol.html#mxnet.symbol.sum) and [log](https://mxnet.incubator.apache.org/api/python/symbol/symbol.html#mxnet.symbol.log) symbols in the process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def diagonal_gaussian_kl(mean: mx.sym.Symbol, std: mx.sym.Symbol) -> mx.sym.Symbol:\n",
    "    \"\"\"\n",
    "    Computes the KL divergence of a standard normal distribution from a diagonal Gaussian.\n",
    "    \n",
    "    :param mean: The mean of the diagonal Gaussian.\n",
    "    :param std: The standard deviations of the diagonal Gaussian.\n",
    "    :return: The value of the Kl divergence.\n",
    "    \"\"\"\n",
    "    # TODO: implement the diagonal gaussian Kl divergence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 2.b\n",
    "We can finally implement the Gaussian reparametrisation. Fill in the correct code to perform the sampling in the sample_latent_state method.\n",
    "\n",
    "**Hint:** In this exercise you will need to draw a random Gaussian sample (see [here](https://mxnet.incubator.apache.org/api/python/symbol/symbol.html#mxnet.symbol.random_normal)). The operator requires a\n",
    "shape whose first entry is the batch size. The batch size is not known to you during implementation, however.\n",
    "You can leave it underspecified by choosing $0$ as a value. When you combine the sampling operator with another\n",
    "operator mxnet will infer the correct the batch size for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GaussianInferenceNetwork(InferenceNetwork):\n",
    "    \"\"\"\n",
    "    An inference network that predicts the parameters of a diagonal Gaussian and samples from that distribution.\n",
    "\n",
    "    :param latent_variable_size: The dimensionality of the latent variable.\n",
    "    :param layer_sizes: Size of each hidden layer in the network.\n",
    "    :param act_type: The activation after each hidden layer.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, latent_variable_size: int, layer_sizes: List[int], act_type: str):\n",
    "        super().__init__(latent_variable_size, layer_sizes, act_type)\n",
    "\n",
    "    def inference(self, data: mx.sym.Symbol) -> Tuple[mx.sym.Symbol, mx.sym.Symbol]:\n",
    "        \"\"\"\n",
    "        Infer the mean and standard deviation.\n",
    "\n",
    "        :param data: A data sample.\n",
    "        :return: The mean and standard deviation.\n",
    "        \"\"\"\n",
    "        # We choose to share the first layer between the networks that compute the standard deviations\n",
    "        # and means. This is a fairly standard design choice.\n",
    "        shared_layer = mx.sym.FullyConnected(data=data, num_hidden=self.layer_sizes[0], name=\"inf_joint_fc\")\n",
    "        shared_layer = mx.sym.Activation(data=shared_layer, act_type=self.act_type, name=\"inf_joint_act\")\n",
    "\n",
    "        prev_out = shared_layer\n",
    "        for i, size in enumerate(self.layer_sizes[1:]):\n",
    "            # TODO: do what you did for the generator in exercise 1.\n",
    "            prev_out = mean_act_i\n",
    "        mean = mx.sym.FullyConnected(data=prev_out, num_hidden=self.latent_var_size, name=\"inf_mean_compute\")\n",
    "\n",
    "        prev_out = shared_layer\n",
    "        for i, size in enumerate(self.layer_sizes[1:]):\n",
    "            # TODO: again, same exercise.\n",
    "            prev_out = var_act_i\n",
    "        # TODO: The standard deviation is constrained to be positive for non-degenerate Gaussians. Take the\n",
    "        # prev_out symbol and transform it into a vector of positive reals.\n",
    "        std = None\n",
    "\n",
    "        return mean, std\n",
    "    \n",
    "    def get_kl_loss(self, mean: mx.sym.Symbol, std: mx.sym.Symbol) -> mx.sym.Symbol:\n",
    "        return mx.sym.MakeLoss(diagonal_gaussian_kl(mean, std), name=\"gaussian_kl_loss\")\n",
    "\n",
    "    def sample_latent_state(self, mean: mx.sym.Symbol, std: mx.sym.Symbol, n: Optional[int] = 1) -> mx.sym.Symbol:\n",
    "        \"\"\"\n",
    "        Sample a value for the latent Gaussian variable.\n",
    "\n",
    "        :param mean: The mean of the Gaussian.\n",
    "        :param std: The standard deviation of the Gaussian.\n",
    "        :param n: Number of samples to be produced.\n",
    "        :return: A Gaussian sample.\n",
    "        \"\"\"\n",
    "        if n > 1:\n",
    "            # TODO: make sure this can be used during training\n",
    "            mean = mx.sym.tile(data=mean, reps=(n,1), name=\"inf_net_replicate_mean\")\n",
    "            std = mx.sym.tile(data=std, reps=(n,1), name=\"inf_net_replicate_std\")\n",
    "        # TODO: This is where the magic happens! Draw a sample from the Gaussian using the Gaussian reparametrisation\n",
    "        # trick and return it.\n",
    "        pass\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exercise 3\n",
    "\n",
    "The only thing that is left to do is to implement VAE training. To train our VAE we will maximise the ELBO:\n",
    "$$ \\mathbb{E}\\left[ \\log p(x|z) \\right] - \\text{KL}(q(z|x)||p(z)) \\ . $$\n",
    "By default, MXNet's optimisers minimise losses instead of maximising objectives. We thus turn the ELBO into a loss by taking its negative. We already get the (sampled) cross entropy loss and (analytically computed) KL divergence as return values form calls to the Generator and inference network respectively. Your job is to combine them into one loss. In mxnet this can be done using the [Group symbol](https://mxnet.incubator.apache.org/api/python/symbol/symbol.html#mxnet.symbol.Group). If a list of losses is provided as an argument to this symbol, the optimiser will add up these losses.\n",
    "\n",
    "**Note**: We could just as well have added up the values of the cross-entropy and the KL divergence and put the resulting loss into a MakeLoss symbol. The reason why we want to separate these two terms of the loss is that it allows us to output KL divergence. We will use this below when defining the KLMetric. Besides tracking the ELBO, which is our optimisation objective, we can also track to KL divergence to see how far the variational approximation moves away from the prior."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GaussianVAE(VAE):\n",
    "    \"\"\"\n",
    "    A VAE with Gaussian latent variables. It assumes a standard normal prior on the latent values.\n",
    "\n",
    "    :param generator: A generator network that specifies the likelihood of the model.\n",
    "    :param inference_net: An inference network that specifies the Gaussian over latent values.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "                 generator: Generator,\n",
    "                 inference_net: GaussianInferenceNetwork,\n",
    "                 kl_divergence: Callable) -> None:\n",
    "        self.generator = generator\n",
    "        self.inference_net = inference_net\n",
    "        self.kl_divergence = kl_divergence\n",
    "\n",
    "    def train(self, data: mx.sym.Symbol, label: mx.sym.Symbol) -> mx.sym.Symbol:\n",
    "        \"\"\"\n",
    "        Train the generator and inference network jointly by optimising the ELBO.\n",
    "\n",
    "        :param data: The training data.\n",
    "        :param label: Copy of the training data.\n",
    "        :return: A list of loss symbols.\n",
    "        \"\"\"\n",
    "        # TODO implement the train method. Make sure to include all relevant loss terms.\n",
    "        pass\n",
    "\n",
    "    def generate_reconstructions(self, data: mx.sym.Symbol, n: int) -> mx.sym.Symbol:\n",
    "        \"\"\"\n",
    "        Generate a number of reconstructions of input data points.\n",
    "\n",
    "        :param data: The input data.\n",
    "        :param n: Number of reconstructions per data point.\n",
    "        :return: The reconstructed data.\n",
    "        \"\"\"\n",
    "        mean, std = self.inference_net.inference(data=data)\n",
    "        mean = mx.sym.tile(data=mean, reps=(n, 1))\n",
    "        std = mx.sym.tile(data=std, reps=(n, 1))\n",
    "        latent_state = self.inference_net.sample_latent_state(mean, std, n)\n",
    "        return self.generator.generate_sample(latent_state=latent_state)\n",
    "\n",
    "    def phantasize(self, n: int) -> mx.sym.Symbol:\n",
    "        \"\"\"\n",
    "        Generate data by randomly sampling from the prior.\n",
    "\n",
    "        :param n: Number of sampled data points.\n",
    "        :return: Randomly generated data points.\n",
    "        \"\"\"\n",
    "        latent_state = mx.sym.random_normal(loc=0, scale=1, shape=(n, self.inference_net.latent_var_size))\n",
    "        return self.generator.generate_sample(latent_state=latent_state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Constructing a VAE\n",
    "\n",
    "We have now all the code for VAEs in place. Below, we have defined a factory method that makes it easier for you to play with different architectures."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def construct_vae(latent_type: str,\n",
    "                  likelihood: str,\n",
    "                  generator_layer_sizes: List[int],\n",
    "                  infer_layer_sizes: List[int],\n",
    "                  latent_variable_size: int,\n",
    "                  data_dims: int,\n",
    "                  generator_act_type: str = \"tanh\",\n",
    "                  infer_act_type: str = \"tanh\") -> VAE:\n",
    "    \"\"\"\n",
    "    Construct a variational autoencoder\n",
    "\n",
    "    :param latent_type: Distribution of latent variable.\n",
    "    :param likelihood: Type of likelihood.\n",
    "    :param generator_layer_sizes: Sizes of generator hidden layers.\n",
    "    :param infer_layer_size: Sizes of inference network hidden layers.\n",
    "    :param latent_variable_size: Size of the latent variable.\n",
    "    :param data_dims: Dimensionality of the data.\n",
    "    :param generator_act_type: Activation function for generator hidden layers.\n",
    "    :param infer_act_type: Activation function for inference network hidden layers.\n",
    "    :return: A variational autoencoder.\n",
    "    \"\"\"\n",
    "    if likelihood == \"bernoulliProd\":\n",
    "        generator = ProductOfBernoullisGenerator(data_dims=data_dims, layer_sizes=generator_layer_sizes,\n",
    "                                                 act_type=generator_act_type)\n",
    "    else:\n",
    "        raise Exception(\"{} is an invalid likelihood type.\".format(likelihood))\n",
    "\n",
    "    if latent_type == \"gaussian\":\n",
    "        inference_net = GaussianInferenceNetwork(latent_variable_size=latent_variable_size,\n",
    "                                                 layer_sizes=infer_layer_sizes,\n",
    "                                                 act_type=infer_act_type)\n",
    "    else:\n",
    "        raise Exception(\"{} is an invalid latent variable type.\".format(latent_type))\n",
    "        \n",
    "    return GaussianVAE(generator=generator, inference_net=inference_net, kl_divergence=diagonal_gaussian_kl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exercise 4\n",
    "\n",
    "Your turn! Construct your own VAE below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae = construct_vae(latent_type=\"gaussian\", likelihood=\"bernoulliProd\", generator_layer_sizes=[200,500],\n",
    "                   infer_layer_sizes=[500,200], latent_variable_size=200, data_dims=784, generator_act_type=\"tanh\",\n",
    "                   infer_act_type=\"tanh\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To check what your VAE looks like we will visualise it. The variables are \"data\" and \"label\" are unbound variables in the computation graph. We will shortly use them to supply data to the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = mx.sym.Variable(\"data\")\n",
    "label = mx.sym.Variable(\"label\")\n",
    "mx.viz.plot_network(vae.train(data, label))\n",
    "# Comment out the line above and uncomment the one below if you want to save the VAE picture on disk\n",
    "#mx.viz.plot_network(vae.train(data, label), title=\"my_vae\", save_format=\"pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training a VAE\n",
    "\n",
    "We are now set to train the VAE. In order to do so we have to define a data iterator and a module in mxnet. The data iterator takes care of batching the data while the module executes the computation graph and updates the model parameters. For the purpose of training the model, we only need the training data which we load from disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist = {}\n",
    "file_name = join(data_dir, \"binary_mnist.{}\".format(TRAIN_SET))\n",
    "logging.info(\"Reading {} into memory\".format(file_name))\n",
    "mnist[TRAIN_SET] = mx.nd.array(np.genfromtxt(file_name))\n",
    "logging.info(\"{} contains {} data points\".format(file_name, mnist[TRAIN_SET].shape[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using that data we define the training set iterator. At this point we also need to choose a batch_size that will then be used in training. Notice that we randomise the order in which the data points are presented in the training set. It is important that `data_name` and `label_name` match the names of the unbound variables in our model. MXNet will use this information to pass the correct data points to the variables. Finally, notice that the data and labels are identical. This is because the VAE tries to reconstruct its input data and thus labels and data are the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 200\n",
    "train_iter = mx.io.NDArrayIter(data=mnist[TRAIN_SET], data_name=\"data\", label=mnist[TRAIN_SET], label_name=\"label\",\n",
    "                                   batch_size=batch_size, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can define a module that will do the training for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae_module = mx.module.Module(vae.train(data=mx.sym.Variable(\"data\"), label=mx.sym.Variable(\"label\")),\n",
    "                           data_names=[train_iter.provide_data[0][0]],\n",
    "                           label_names=[train_iter.provide_label[0][0]], context=ctx, logger=logging)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training the vae (or any other network defined in mxnet) is most easily done using the [fit method](https://mxnet.incubator.apache.org/api/python/module/module.html#mxnet.module.BaseModule.fit). Here we choose to train our model for 20 epochs. Our optimiser will be adam. Training will take some time (approximately 2-5 minutes, depending on your machine). We are keeping track of the loss (the negative ELBO) to see how the model develops. Notice that this loss\n",
    "is not the actual negative ELBO but a _doubly stochastic approximation_ to it (see [here](https://projecteuclid.org/download/pdf_1/euclid.aoms/1177729586) and [here](http://proceedings.mlr.press/v32/titsias14.pdf)). The two sources of stochasticity are the mini-batches (the ELBO is defined with respect to the entire data set) and the Gaussian reparametrisation (we approximate the integral over $ z $ through sampling). Both sources of stochasticity leave the approximation unbiased, however. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ElboMetric(mx.metric.EvalMetric):\n",
    "\n",
    "    def __init__(self, name: str = \"elbo\",\n",
    "                 output_names: Optional[List[str]] = None,\n",
    "                 label_names: Optional[List[str]] = None):\n",
    "        super().__init__(name, output_names, label_names)\n",
    "\n",
    "    def update(self, labels: List[mx.nd.array], preds: List[mx.nd.array]):\n",
    "        neg_likelihoods, kl_values, *_ = preds\n",
    "        label = labels[0]\n",
    "        batch_size = label.shape[0]\n",
    "        self.num_inst += batch_size\n",
    "        self.sum_metric += -mx.nd.sum(neg_likelihoods).asscalar() - mx.nd.sum(kl_values).asscalar()\n",
    "\n",
    "\n",
    "class KLMetric(mx.metric.EvalMetric):\n",
    "\n",
    "    def __init__(self, name: str = \"kl_divergence\",\n",
    "                 output_names: Optional[List[str]] = None,\n",
    "                 label_names: Optional[List[str]] = None):\n",
    "        super().__init__(name, output_names, label_names)\n",
    "\n",
    "    def update(self, labels: List[mx.nd.array], preds: List[mx.nd.array]):\n",
    "        neg_likelihoods, kl_values, *_ = preds\n",
    "        label = labels[0]\n",
    "        batch_size = label.shape[0]\n",
    "        self.num_inst += batch_size\n",
    "        self.sum_metric += mx.nd.sum(kl_values).asscalar()\n",
    "        \n",
    "metric = mx.metric.CompositeEvalMetric([ElboMetric(), KLMetric()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 20\n",
    "optimiser = \"adam\"\n",
    "\n",
    "vae_module.fit(train_data=train_iter, optimizer=optimiser, force_init=True, force_rebind=True, num_epoch=epochs,\n",
    "               optimizer_params={'learning_rate': DEFAULT_LEARNING_RATE},\n",
    "               batch_end_callback=mx.callback.Speedometer(frequent=20, batch_size=batch_size),\n",
    "               epoch_end_callback=mx.callback.do_checkpoint('vae'),\n",
    "               eval_metric=metric)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The VAE has converged nicely (arguably we could have trained it even longer). Observe how the KL divergence has increased towards the end of training. This means that the variational approximation has become increasingly specialised and distinct from the standard normal prior."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simulating Data from the VAE\n",
    "\n",
    "Now that we have trained the VAE, we can use it to produce data for us. We first need to load the parameters and supply them to a new symbol that does the phantasizing. We then produce 10 random digits. The pixel intensities are the predicted Bernoulli parameters. **NOTE**: The width and height variables were computed in the beginning when we were first looking at our data set. We are simply reusing them here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_samples = 10\n",
    "\n",
    "params = vae_module.get_params()[0]\n",
    "# We group the outputs of phantasize to be able to process them as a single symbol\n",
    "dream_digits = mx.sym.Group(vae.phantasize(num_samples))\n",
    "# construct an executor by binding the parameters to the learned values\n",
    "dream_exec = dream_digits.bind(ctx=ctx, args=params)\n",
    "\n",
    "# run the computation\n",
    "digits, latent_values = dream_exec.forward()\n",
    "\n",
    "# transform output into numpy arrays\n",
    "digits = digits.asnumpy()\n",
    "latent_values = latent_values.asnumpy()\n",
    "rows = int(num_samples / 5)\n",
    "\n",
    "plot, axes = plt.subplots(rows, 5,  sharex='col', sharey='row', figsize=(30,6))\n",
    "sample=0\n",
    "for row in range(rows):\n",
    "    for col in range(5):\n",
    "        axes[row][col].imshow(np.reshape(digits[sample,:],(width,height)), cmap=cm.Greys)\n",
    "        sample += 1\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that the above samples were obtained by sampling from the standard normal prior. We can also check the reconstruction abilities of the VAE. To do this we sample a random number from the test set and generate 10 reconstructions of it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_name = join(data_dir, \"binary_mnist.{}\".format(TEST_SET))\n",
    "logging.info(\"Reading {} into memory\".format(file_name))\n",
    "test_set = np.genfromtxt(file_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is what the sampled number looks like. (Rerun if you want a different number.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random_idx = np.random.randint(test_set.shape[0])\n",
    "random_picture = test_set[random_idx, :]\n",
    "plot, canvas = plt.subplots(1, figsize=(5,5))\n",
    "canvas.imshow(np.reshape(random_picture,(width,height)), cmap=cm.Greys)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now we are going to see how the VAE reconstructs it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_samples = 10\n",
    "\n",
    "params = vae_module.get_params()[0]\n",
    "# We group the outputs of generate_reconstructions to be able to process them as a single symbol\n",
    "reconstructions = mx.sym.Group(vae.generate_reconstructions(mx.sym.Variable(\"random_digit\"), num_samples))\n",
    "# construct an executor by binding the parameters to the learned values\n",
    "params[\"random_digit\"] = mx.nd.array(random_picture.reshape((1,height*width)))\n",
    "reconstruction_exec = reconstructions.bind(ctx=ctx, args=params)\n",
    "\n",
    "# run the computation\n",
    "digits, latent_values = reconstruction_exec.forward()\n",
    "\n",
    "# transform output into numpy arrays\n",
    "digits = digits.asnumpy()\n",
    "latent_values = latent_values.asnumpy()\n",
    "\n",
    "# plot the reconstructed digits\n",
    "rows = int(num_samples / 5)\n",
    "plot, axes = plt.subplots(rows, 5,  sharex='col', sharey='row', figsize=(30,6))\n",
    "sample=0\n",
    "for row in range(rows):\n",
    "    for col in range(5):\n",
    "        axes[row][col].imshow(np.reshape(digits[sample,:],(width,height)), cmap=cm.Greys)\n",
    "        sample += 1\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Congrats!\n",
    "\n",
    "You have completed your first VAE tutorial. From here on you can modify this code to play with your own VAE ideas. Want to you use different data? Or vary the neural net structure for the generator or inference networks? Or use a different likelihood model? The world of probablistic modeling is yours to explore."
   ]
  }
 ],
 "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}