VAE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Teaching a Variational Autoencoder (VAE) to draw MNIST characters\n",
    "\n",
    "Autoencoders are a type of neural network that can be used to learn efficient codings of input data. \n",
    "Given some inputs, the network firstly applies a series of transformations that map the input data into a lower dimensional space. This part of the network is called the _encoder_. Then, the network uses the encoded data to try and recreate the inputs. This part of the network is the _decoder_. Using the encoder, we can later compress data of the type that is understood by the network. However, autoencoders are rarely used for this purpose, as usually there exist hand-crafted algorithms (like _jpg_-compression) that are more efficient. Instead, autoencoders have repeatedly been applied to perform denoising tasks. Then, the encoder receives pictures that have been tampered with noise, and it learns how to reconstruct the original images.\n",
    "\n",
    "\n",
    "## Variational Autoencoders put simply\n",
    "But there exists a much more interesting application for autoencoders. This application is called the _variational autoencoder_. Using variational autoencoders, it's not only possible to compress data -- it's also possible to generate new objects of the type the autoencoder has seen before.\n",
    "\n",
    "Using a general autoencoder, we don't know anything about the coding that's been generated by our network. We could take a look at and compare different encoded objects, but it's unlikely that we'll be able to understand what's going on. This means that we won't be able to use our decoder for creating new objects -- we simply don't know what the inputs should look like.\n",
    "\n",
    "Using a variational autoencoder, we take the opposite approach instead. We will not try to make guesses concerning the distribution that's being followed by the latent vectors. We simply tell our network what we want this distribution to look like. Usually, we will constrain the network to produce latent vectors having entries that follow the unit normal distribution. Then, when trying to generate data, we can simply sample some values from this distribution, feed them to the decoder, and the decoder will return us completely new objects that appear just like the objects our network has been trained with.\n",
    "\n",
    "Let's see how this can be done using python and tensorflow. We are going to teach our network how to draw MNIST characters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## First steps -- Loading the training data\n",
    "Firstly, we perform some basic imports. Tensorflow has a quite handy function that allows us to easily access the MNIST data set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "from tensorflow.examples.tutorials.mnist import input_data\n",
    "mnist = input_data.read_data_sets('MNIST_data')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining our input and output data\n",
    "MNIST images have a dimension of 28 * 28 pixels with one color channel. Our inputs _X_in_ will be batches of MNIST characters, while our network will learn to reconstruct them and output them in a placeholder _Y_, which thus has the same dimensions. _Y_flat_ will be used later, when computing losses. _keep_prob_ will be used when applying dropouts as a means of regularization. During training, it will have a value of 0.8. When generating new data, we won't apply dropout, so the value will be 1. The function _lrelu_ is being defined as tensorflow unfortunately doesn't come up with a predefined leaky ReLU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "\n",
    "batch_size = 64\n",
    "\n",
    "X_in = tf.placeholder(dtype=tf.float32, shape=[None, 28, 28], name='X')\n",
    "Y    = tf.placeholder(dtype=tf.float32, shape=[None, 28, 28], name='Y')\n",
    "Y_flat = tf.reshape(Y, shape=[-1, 28 * 28])\n",
    "keep_prob = tf.placeholder(dtype=tf.float32, shape=(), name='keep_prob')\n",
    "\n",
    "dec_in_channels = 1\n",
    "n_latent = 8\n",
    "\n",
    "reshaped_dim = [-1, 7, 7, dec_in_channels]\n",
    "inputs_decoder = 49 * dec_in_channels / 2\n",
    "\n",
    "\n",
    "def lrelu(x, alpha=0.3):\n",
    "    return tf.maximum(x, tf.multiply(x, alpha))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining the encoder\n",
    "As our inputs are images, it's most reasonable to apply some convolutional transformations to them. What's most noteworthy is the fact that we are creating two vectors in our encoder, as the encoder is supposed to create objects following a Gaussian Distribution:\n",
    "* A vector of means\n",
    "* A vector of standard deviations\n",
    "\n",
    "You will see later how we \"force\" the encoder to make sure it really creates values following a Normal Distribution. The returned values that will be fed to the decoder are the _z_-values. We will need the mean and standard deviation of our distributions later, when computing losses. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encoder(X_in, keep_prob):\n",
    "    activation = lrelu\n",
    "    with tf.variable_scope(\"encoder\", reuse=None):\n",
    "        X = tf.reshape(X_in, shape=[-1, 28, 28, 1])\n",
    "        x = tf.layers.conv2d(X, filters=64, kernel_size=4, strides=2, padding='same', activation=activation)\n",
    "        x = tf.nn.dropout(x, keep_prob)\n",
    "        x = tf.layers.conv2d(x, filters=64, kernel_size=4, strides=2, padding='same', activation=activation)\n",
    "        x = tf.nn.dropout(x, keep_prob)\n",
    "        x = tf.layers.conv2d(x, filters=64, kernel_size=4, strides=1, padding='same', activation=activation)\n",
    "        x = tf.nn.dropout(x, keep_prob)\n",
    "        x = tf.contrib.layers.flatten(x)\n",
    "        mn = tf.layers.dense(x, units=n_latent)\n",
    "        sd       = 0.5 * tf.layers.dense(x, units=n_latent)            \n",
    "        epsilon = tf.random_normal(tf.stack([tf.shape(x)[0], n_latent])) \n",
    "        z  = mn + tf.multiply(epsilon, tf.exp(sd))\n",
    "        \n",
    "        return z, mn, sd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining the decoder\n",
    "The decoder does not care about whether the input values are sampled from some specific distribution that has been defined by us. It simply will try to reconstruct the input images. To this end, we use a series of transpose convolutions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def decoder(sampled_z, keep_prob):\n",
    "    with tf.variable_scope(\"decoder\", reuse=None):\n",
    "        x = tf.layers.dense(sampled_z, units=inputs_decoder, activation=lrelu)\n",
    "        x = tf.layers.dense(x, units=inputs_decoder * 2 + 1, activation=lrelu)\n",
    "        x = tf.reshape(x, reshaped_dim)\n",
    "        x = tf.layers.conv2d_transpose(x, filters=64, kernel_size=4, strides=2, padding='same', activation=tf.nn.relu)\n",
    "        x = tf.nn.dropout(x, keep_prob)\n",
    "        x = tf.layers.conv2d_transpose(x, filters=64, kernel_size=4, strides=1, padding='same', activation=tf.nn.relu)\n",
    "        x = tf.nn.dropout(x, keep_prob)\n",
    "        x = tf.layers.conv2d_transpose(x, filters=64, kernel_size=4, strides=1, padding='same', activation=tf.nn.relu)\n",
    "        \n",
    "        x = tf.contrib.layers.flatten(x)\n",
    "        x = tf.layers.dense(x, units=28*28, activation=tf.nn.sigmoid)\n",
    "        img = tf.reshape(x, shape=[-1, 28, 28])\n",
    "        return img"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we'll wire together both parts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampled, mn, sd = encoder(X_in, keep_prob)\n",
    "dec = decoder(sampled, keep_prob)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Computing losses and enforcing a Gaussian latent distribution\n",
    "For computing the image reconstruction loss, we simply use squared difference (which could lead to images sometimes looking a bit fuzzy). This loss is combined with the _Kullback-Leibler divergence_, which makes sure our latent values will be sampled from a normal distribution. For more on this topic, please take a look a Jaan Altosaar's great article on VAEs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unreshaped = tf.reshape(dec, [-1, 28*28])\n",
    "img_loss = tf.reduce_sum(tf.squared_difference(unreshaped, Y_flat), 1)\n",
    "latent_loss = -0.5 * tf.reduce_sum(1.0 + 2.0 * sd - tf.square(mn) - tf.exp(2.0 * sd), 1)\n",
    "loss = tf.reduce_mean(img_loss + latent_loss)\n",
    "optimizer = tf.train.AdamOptimizer(0.0005).minimize(loss)\n",
    "sess = tf.Session()\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Training the network\n",
    "Now, we can finally train our VAE! Every 200 steps, we'll take a look at what the current reconstructions look like. After having processed about 2000 batches, most reconstructions will look reasonable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(30000):\n",
    "    batch = [np.reshape(b, [28, 28]) for b in mnist.train.next_batch(batch_size=batch_size)[0]]\n",
    "    sess.run(optimizer, feed_dict = {X_in: batch, Y: batch, keep_prob: 0.8})\n",
    "        \n",
    "    if not i % 200:\n",
    "        ls, d, i_ls, d_ls, mu, sigm = sess.run([loss, dec, img_loss, latent_loss, mn, sd], feed_dict = {X_in: batch, Y: batch, keep_prob: 1.0})\n",
    "        plt.imshow(np.reshape(batch[0], [28, 28]), cmap='gray')\n",
    "        plt.show()\n",
    "        plt.imshow(d[0], cmap='gray')\n",
    "        plt.show()\n",
    "        print(i, ls, np.mean(i_ls), np.mean(d_ls))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating new data\n",
    "The most awesome part is that we are now able to create new characters. To this end, we simply sample values from a unit normal distribution and feed them to our decoder. Most of the created characters look just like they've been written by humans.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "randoms = [np.random.normal(0, 1, n_latent) for _ in range(10)]\n",
    "imgs = sess.run(dec, feed_dict = {sampled: randoms, keep_prob: 1.0})\n",
    "imgs = [np.reshape(imgs[i], [28, 28]) for i in range(len(imgs))]\n",
    "\n",
    "for img in imgs:\n",
    "    plt.figure(figsize=(1,1))\n",
    "    plt.axis('off')\n",
    "    plt.imshow(img, cmap='gray')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Conclusion\n",
    "Now, this obviously is a relatively simple example of an application of VAEs. But just think about what could be possible! Neural networks could learn to compose music. They could automatically create illustrations for books, games etc. With a bit of creativity, VAEs will open up space for some awesome projects "
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}