nn-theano.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Speeding up your Neural Network with Theano and the GPU\n",
    "\n",
    "In a [previous blog post](http://www.wildml.com/2015/09/implementing-a-neural-network-from-scratch/) we build a simple Neural Network from scratch. Let's build on top of this and speed up our code using the [Theano](http://deeplearning.net/software/theano/) library. With Theano we can make our code not only faster, but also more concise!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is Theano?\n",
    "\n",
    "Theano describes itself as a Python library that lets you to define, optimize, and evaluate mathematical expressions, especially ones with multi-dimensional arrays. The way I understand Theano is that it allows me to define **graphs of computations**. Under the hood Theano optimizes these computations in a [variety of ways](http://deeplearning.net/software/theano/introduction.html#introduction), including avoiding redundant calculations, generating optimized C code, and (optionally) using the GPU. Theano also has the capability to automatically [differentiate](http://deeplearning.net/software/theano/tutorial/gradients.html) mathematical expressions. By modeling computations as graphs it can calculate complex gradients using the chain rule. This means we no longer need to compute the gradients ourselves!\n",
    "\n",
    "Because Neural Networks are easily expressed as graphs of computations, Theano is a great fit. It's probably the main use case and Theano includes several [convenience functions for neural networks](http://deeplearning.net/software/theano/library/tensor/nnet/index.html#module-nnet)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Setup\n",
    "\n",
    "The setup is identical to that in [Implementing a Neural Network from Scratch](http://www.wildml.com/2015/09/implementing-a-neural-network-from-scratch/), which I recommend you read (or at least skim) first. I'll just quickly recap: We have two classes (red and blue) and want to train a Neural Network classifier that separates the two. We will train a 3-layer Neural Network, with input layer size 2, output layer size 2, and hidden layer size 3. We will use batch gradient descent with a fixed learning rate to train our network. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import sklearn\n",
    "import sklearn.datasets\n",
    "import matplotlib\n",
    "import theano\n",
    "import theano.tensor as T\n",
    "import pydot\n",
    "from IPython.display import Image\n",
    "from IPython.display import SVG\n",
    "import timeit\n",
    "\n",
    "# Display plots inline and change default figure size\n",
    "%matplotlib inline\n",
    "matplotlib.rcParams['figure.figsize'] = (10.0, 8.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Generate a dataset and plot it\n",
    "np.random.seed(0)\n",
    "train_X, train_y = sklearn.datasets.make_moons(200, noise=0.20)\n",
    "train_X = train_X.astype(np.float32)\n",
    "train_y = train_y.astype(np.int32)\n",
    "plt.scatter(train_X[:,0], train_X[:,1], s=40, c=train_y, cmap=plt.cm.Spectral)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Helper function to plot a decision boundary.\n",
    "# If you don't fully understand this function don't worry, it just generates the contour plot.\n",
    "def plot_decision_boundary(pred_func):\n",
    "    # Set min and max values and give it some padding\n",
    "    x_min, x_max = train_X[:, 0].min() - .5, train_X[:, 0].max() + .5\n",
    "    y_min, y_max = train_X[:, 1].min() - .5, train_X[:, 1].max() + .5\n",
    "    h = 0.01\n",
    "    # Generate a grid of points with distance h between them\n",
    "    xx, yy = np.meshgrid(np.arange(x_min, x_max, h), np.arange(y_min, y_max, h))\n",
    "    # Predict the function value for the whole gid\n",
    "    Z = pred_func(np.c_[xx.ravel(), yy.ravel()])\n",
    "    Z = Z.reshape(xx.shape)\n",
    "    # Plot the contour and training examples\n",
    "    plt.contourf(xx, yy, Z, cmap=plt.cm.Spectral)\n",
    "    plt.scatter(train_X[:, 0], train_X[:, 1], c=train_y, cmap=plt.cm.Spectral)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Size definitions\n",
    "num_examples = len(train_X) # training set size\n",
    "nn_input_dim = 2 # input layer dimensionality\n",
    "nn_output_dim = 2 # output layer dimensionality\n",
    "nn_hdim = 100 # hiden layer dimensionality\n",
    "\n",
    "# Gradient descent parameters (I picked these by hand)\n",
    "epsilon = 0.01 # learning rate for gradient descent\n",
    "reg_lambda = 0.01 # regularization strength "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining the Computation Graph in Theano\n",
    "\n",
    "The first thing we need to is define our computations using Theano. We start by defining our input data matrix `X` and our training labels `y`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Our data vectors\n",
    "X = T.matrix('X') # matrix of doubles\n",
    "y = T.lvector('y') # vector of int64"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's crucial thing to understand: We have not assigned any values to `X` or `y`. All we have done is defined mathematical expressions for them. We can use these expressions in subsequent calculations. If we want to evaluate an expression we can call its `eval` method. For example, to evaluate the expression `X * 2` for a given value of `X` we could do the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "(X * 2).eval({X : [[1,1],[2,2]] })"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Theano handles the type checking for us, which is very useful when defining more complex expressions. Trying to assign a value of the wrong data type to `X` would result in an error. [Here is the full list of Theano types](http://deeplearning.net/software/theano/library/tensor/basic.html#all-fully-typed-constructors).\n",
    "\n",
    "`X` and `y` above are *stateless*. Whenever we want to evaluate an expression that depends on them we need to provide their values. Theano also has something called [shared variables](http://deeplearning.net/software/theano/tutorial/examples.html#using-shared-variables), which have internal state associated with them. Their value that is kept in memory and can be shared by all functions that use them. Shared variables can also be updated, and Theano includes low-level optimizations that makes updating them very efficient, especially on GPUs. Our network parameters $W_1, b_1, W_2, b_2$ are constantly updated using gradient descent, so they are perfect candidates for shared variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Shared variables with initial values. We need to learn these.\n",
    "W1 = theano.shared(np.random.randn(nn_input_dim, nn_hdim), name='W1')\n",
    "b1 = theano.shared(np.zeros(nn_hdim), name='b1')\n",
    "W2 = theano.shared(np.random.randn(nn_hdim, nn_output_dim), name='W2')\n",
    "b2 = theano.shared(np.zeros(nn_output_dim), name='b2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's define expressions for our forward propagation. The calculations are identical to what we did in our [previous implementation](http://www.wildml.com/2015/09/implementing-a-neural-network-from-scratch/), just that we are defining Theano expressions. Again, remember that these expressions are not evaluated, we are just defining them. You can think of them as lambda expressions that require input values when called. We also use some of Theano's convenience functions like [nnet.softmax](http://deeplearning.net/software/theano/library/tensor/nnet/nnet.html#tensor.nnet.softmax) and [nnet.categorical_crossentropy](http://deeplearning.net/software/theano/library/tensor/nnet/nnet.html#tensor.nnet.categorical_crossentropy) replace our manual implementations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Forward propagation\n",
    "# Note: We are just defining the expressions, nothing is evaluated here!\n",
    "z1 = X.dot(W1) + b1\n",
    "a1 = T.tanh(z1)\n",
    "z2 = a1.dot(W2) + b2\n",
    "y_hat = T.nnet.softmax(z2) # output probabilties\n",
    "\n",
    "# The regularization term (optional)\n",
    "loss_reg = 1./num_examples * reg_lambda/2 * (T.sum(T.sqr(W1)) + T.sum(T.sqr(W2))) \n",
    "# the loss function we want to optimize\n",
    "loss = T.nnet.categorical_crossentropy(y_hat, y).mean() + loss_reg\n",
    "\n",
    "# Returns a class prediction\n",
    "prediction = T.argmax(y_hat, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We saw how we can evaluate a Theano expression by calling its `eval` method. A much more convenient way is to create a [Theano function](http://deeplearning.net/software/theano/library/compile/function.html#function.function) for expressions we want to evaluate. To create a function we need to define its inputs and outputs. For example, to calculate the loss, we need to know the values for $X$ and $y$. Once created, we can call it function just like any other Python function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Theano functions that can be called from our Python code\n",
    "forward_prop = theano.function([X], y_hat)\n",
    "calculate_loss = theano.function([X, y], loss)\n",
    "predict = theano.function([X], prediction)\n",
    "\n",
    "# Example call: Forward Propagation\n",
    "forward_prop([[1,2]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now is a good time to get a sense of how Theano constructs a computational graph. Looking at the expressions for $\\hat{y}$, we can see that it depends on $z2$, which in turn depends on $a_1$, $W_2$ and $b_2$, and so on. Theano lets us visualize this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "theano.printing.pydotprint(forward_prop, var_with_name_simple=True, compact=True, outfile='img/nn-theano-forward_prop.png', format='png')\n",
    "SVG(theano.printing.pydotprint(forward_prop, var_with_name_simple=True, compact=True, return_image=True, format='svg'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the *optimized* computational graph that Theano has constructed for our `forward_prop` function. We can also get a textual description:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "theano.printing.debugprint(forward_prop)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What's left is defining the updates to the network parameters we use with gradient descent. We previously calculated the gradients using backpropagation. We could express the same calculations using Theano (see code that's commented out below), but it's much easier if we let Theano calculate the derivatives for us! We need the derivates of our loss function $L$ with respect to our parameters: $\\frac{\\partial L}{\\partial W_2}$, $\\frac{\\partial L}{\\partial b_2}$, $\\frac{\\partial L}{\\partial W_1}$, $\\frac{\\partial L}{\\partial b_1}$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Easy: Let Theano calculate the derivatives for us!\n",
    "dW2 = T.grad(loss, W2)\n",
    "db2 = T.grad(loss, b2)\n",
    "dW1 = T.grad(loss, W1)\n",
    "db1 = T.grad(loss, b1)\n",
    "\n",
    "# Backpropagation (Manual)\n",
    "# Note: We are just defining the expressions, nothing is evaluated here!\n",
    "# y_onehot = T.eye(2)[y]\n",
    "# delta3 = y_hat - y_onehot\n",
    "# dW2 = (a1.T).dot(delta3) * (1. + reg_lambda)\n",
    "# db2 = T.sum(delta3, axis=0)\n",
    "# delta2 = delta3.dot(W2.T) * (1. - T.sqr(a1))\n",
    "# dW1 = T.dot(X.T, delta2) * (1 + reg_lambda)\n",
    "# db1 = T.sum(delta2, axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because we defined $W_2, b_2, W_1, b_1$ as shared variables we can use Theano's update mechanism to update their values. The following function (without return value) does a single gradient descent update given $X$ and $y$ as inputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "gradient_step = theano.function(\n",
    "    [X, y],\n",
    "    updates=((W2, W2 - epsilon * dW2),\n",
    "             (W1, W1 - epsilon * dW1),\n",
    "             (b2, b2 - epsilon * db2),\n",
    "             (b1, b1 - epsilon * db1)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we don't need to explicitly do a forward propagation here. Theano knows that our gradients depend on our predictions from the forward propagation and it will handle all the necessary calculations for us. It does everything it needs to update the values."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now define a function to train a Neural Network using gradient descent. Again, it's equivalent to what we had in our original code, only that we are now calling the `gradient_step` function defined above instead of doing the calculations ourselves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# This function learns parameters for the neural network and returns the model.\n",
    "# - num_passes: Number of passes through the training data for gradient descent\n",
    "# - print_loss: If True, print the loss every 1000 iterations\n",
    "def build_model(num_passes=20000, print_loss=False):\n",
    "    \n",
    "    # Re-Initialize the parameters to random values. We need to learn these.\n",
    "    # (Needed in case we call this function multiple times)\n",
    "    np.random.seed(0)\n",
    "    W1.set_value(np.random.randn(nn_input_dim, nn_hdim) / np.sqrt(nn_input_dim))\n",
    "    b1.set_value(np.zeros(nn_hdim))\n",
    "    W2.set_value(np.random.randn(nn_hdim, nn_output_dim) / np.sqrt(nn_hdim))\n",
    "    b2.set_value(np.zeros(nn_output_dim))\n",
    "    \n",
    "    # Gradient descent. For each batch...\n",
    "    for i in xrange(0, num_passes):\n",
    "        # This will update our parameters W2, b2, W1 and b1!\n",
    "        gradient_step(train_X, train_y)\n",
    "        \n",
    "        # Optionally print the loss.\n",
    "        # This is expensive because it uses the whole dataset, so we don't want to do it too often.\n",
    "        if print_loss and i % 1000 == 0:\n",
    "            print \"Loss after iteration %i: %f\" %(i, calculate_loss(train_X, train_y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Build a model with a 3-dimensional hidden layer\n",
    "build_model(print_loss=True)\n",
    "\n",
    "# Plot the decision boundary\n",
    "plot_decision_boundary(lambda x: predict(x))\n",
    "plt.title(\"Decision Boundary for hidden layer size 3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it! We've just ported our code over to Theano. I got a 2-3x speedup on my Macbook (it would likely be more if we had larger matrix multiplications). Note that we're not using a GPU yet. Let's do that next!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}