nn-theano-gpu.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running on a GPU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To speed things up further Theano can [make use of your GPU](http://deeplearning.net/software/theano/tutorial/using_gpu.html). Theano relies on Nvidia's CUDA platform, so yo need a [CUDA-enabled graphics card](http://docs.nvidia.com/cuda/cuda-getting-started-guide-for-mac-os-x/#cuda-enabled-gpu). If you don't have one (for example if you're on a Macbook with an integrated GPU like me) you could spin up a [GPU-optimized Amazon EC2 instance](https://aws.amazon.com/ec2/instance-types/#gpu) and try things there. To use the GPU with Theano you must have the CUDA toolkit installed and [set a few environment variables](http://deeplearning.net/software/theano/install.html#gpu-linux). The [repository readme](https://github.com/dennybritz/nn-theano) contains instructions on how to setup CUDA on an EC2 isntance.\n",
    "\n",
    "While any Theano code will run just fine with the GPU enabled you probably want to spend time optimizing your code get optimal performance. Running code that isn't optimized for a GPU may actually be slower on a GPU than a CPU due to excessive data transfers.\n",
    "\n",
    "Here are  some things you must consider when writing Theano code for a GPU:\n",
    "\n",
    "- The default floating point data type is `float64`, but in order to use the GPU you must use `float32`. There are several ways to do this in Theano. You can either use the [fulled typed tensor constructors](http://deeplearning.net/software/theano/library/tensor/basic.html#all-fully-typed-constructors) like `T.fvector` or set  `theano.config.floatX` to `float32` and use the [simple tensor constructors](http://deeplearning.net/software/theano/library/tensor/basic.html#constructors-with-optional-dtype) like `T.vector`.\n",
    "- Shared variables must be initialized with values of type `float32` in order to be stored on the GPU. In other words, when creating a shared variable from a numpy array you must initialize the array with the `dtype=float32` argument, or cast it using the `astype` function. See the [numpy data type documentation](http://docs.scipy.org/doc/numpy/user/basics.types.html) for more details.\n",
    "- Be very careful with data transfers to and from the GPU. Ideally you want all your data on the GPU. Ways to achieve this include initializing commonly used data as shared variables with a `float32` data type, and to [avoid copying return values](http://deeplearning.net/software/theano/tutorial/using_gpu.html#returning-a-handle-to-device-allocated-data) back to the host using the `gpu_from_host` function. Theano includes [profiling capabilities](http://deeplearning.net/software/theano/tutorial/profiling.html) that allow you to debug how much time you are spending on data transfers.\n",
    "\n",
    "With all that in mind, let's modify our code to run well on a GPU. I'll highlight the most important changes we here, but the full code is available [here](https://github.com/dennybritz/nn-theano) or at the bottom of the post."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import sklearn\n",
    "import sklearn.datasets\n",
    "import theano\n",
    "import theano.tensor as T\n",
    "import time\n",
    "\n",
    "# Use float32 as the default float data type\n",
    "theano.config.floatX = 'float32'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, to see a significant speedup we should increase the size of our data so that the computation time will be dominated by matrix multiplication and other mathematical operations (as opposed to looping in Python code and assigning values). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "train_X, train_y = sklearn.datasets.make_moons(5000, noise=0.20)\n",
    "train_X = train_X.astype(np.float32)\n",
    "train_y = train_y.astype(np.int32)\n",
    "nn_hdim = 1000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we change our scalar values to `float32` to ensure that our calculations stay within that type. If we multiplied a `float32` value by a `float64` value the return value would be a `float64`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "epsilon = np.float32(0.01) # learning rate for gradient descent\n",
    "reg_lambda = np.float32(0.01) # regularization strength "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To force the storage of our data on the GPU we use shared variables of type `float32`. Because we are constantly using the input data `X` and `y` to perform gradient updates we store those on the GPU as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# These float32 values will be initialized on the GPU\n",
    "X = theano.shared(train_X.astype('float32'), name='X')\n",
    "y = theano.shared(train_y_onehot.astype('float32'), name='y')\n",
    "W1 = theano.shared(np.random.randn(nn_input_dim, nn_hdim).astype('float32'), name='W1')\n",
    "b1 = theano.shared(np.zeros(nn_hdim).astype('float32'), name='b1')\n",
    "W2 = theano.shared(np.random.randn(nn_hdim, nn_output_dim).astype('float32'), name='W2')\n",
    "b2 = theano.shared(np.zeros(nn_output_dim).astype('float32'), name='b2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We remove the inputs from our functions since the X and y values are stored as shared variables (on the GPU) now. We no longer need to pass them to the function calls."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "gradient_step = theano.function(\n",
    "    inputs=[],\n",
    "    updates=((W2, W2 - epsilon * dW2),\n",
    "             (W1, W1 - epsilon * dW1),\n",
    "             (b2, b2 - epsilon * db2),\n",
    "             (b1, b1 - epsilon * db1)))\n",
    "\n",
    "# Does a gradient step and updates the values of our paremeters\n",
    "gradient_step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running and Evaluation\n",
    "\n",
    "To use a GPU you must run the code with the `THEANO_FLAGS=device=gpu,floatX=float32` environment variable set. I tested the GPU-optimized code on a `g2.2xlarge` AWS EC2 instance. Running one `gradient_step()` on the CPU took around **250ms**. With the GPU enabled it merely took **7.5ms**. That's a 30x speedup, and if our dataset or parameter space were larger we would probably see an even more significant speedup.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%timeit gradient_step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Full GPU-optimzied code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Generate a dataset\n",
    "np.random.seed(0)\n",
    "train_X, train_y = sklearn.datasets.make_moons(5000, noise=0.20)\n",
    "train_y_onehot = np.eye(2)[train_y]\n",
    "\n",
    "# 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 = 1000 # hiden layer dimensionality\n",
    "\n",
    "# Gradient descent parameters (I picked these by hand)\n",
    "epsilon = np.float32(0.01) # learning rate for gradient descent\n",
    "reg_lambda = np.float32(0.01) # regularization strength \n",
    "\n",
    "# GPU NOTE: Conversion to float32 to store them on the GPU!\n",
    "X = theano.shared(train_X.astype('float32')) # initialized on the GPU\n",
    "y = theano.shared(train_y_onehot.astype('float32'))\n",
    "\n",
    "# GPU NOTE: Conversion to float32 to store them on the GPU!\n",
    "W1 = theano.shared(np.random.randn(nn_input_dim, nn_hdim).astype('float32'), name='W1')\n",
    "b1 = theano.shared(np.zeros(nn_hdim).astype('float32'), name='b1')\n",
    "W2 = theano.shared(np.random.randn(nn_hdim, nn_output_dim).astype('float32'), name='W2')\n",
    "b2 = theano.shared(np.zeros(nn_output_dim).astype('float32'), name='b2')\n",
    "\n",
    "# Forward propagation\n",
    "z1 = X.dot(W1) + b1\n",
    "a1 = T.tanh(z1)\n",
    "z2 = a1.dot(W2) + b2\n",
    "y_hat = T.nnet.softmax(z2)\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",
    "# Returns a class prediction\n",
    "prediction = T.argmax(y_hat, axis=1)\n",
    "\n",
    "# Gradients\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",
    "# Note that we removed the input values because we will always use the same shared variable\n",
    "# GPU NOTE: Removed the input values to avoid copying data to the GPU.\n",
    "forward_prop = theano.function([], y_hat)\n",
    "calculate_loss = theano.function([], loss)\n",
    "predict = theano.function([], prediction)\n",
    "\n",
    "# GPU NOTE: Removed the input values to avoid copying data to the GPU.\n",
    "gradient_step = theano.function(\n",
    "    [],\n",
    "    # profile=True,\n",
    "    updates=((W2, W2 - epsilon * dW2),\n",
    "             (W1, W1 - epsilon * dW1),\n",
    "             (b2, b2 - epsilon * db2),\n",
    "             (b1, b1 - epsilon * db1)))\n",
    "\n",
    "def build_model(num_passes=20000, print_loss=False):\n",
    "    # Re-Initialize the parameters to random values. We need to learn these.\n",
    "    np.random.seed(0)\n",
    "    # GPU NOTE: Conversion to float32 to store them on the GPU!\n",
    "    W1.set_value((np.random.randn(nn_input_dim, nn_hdim) / np.sqrt(nn_input_dim)).astype('float32'))\n",
    "    b1.set_value(np.zeros(nn_hdim).astype('float32'))\n",
    "    W2.set_value((np.random.randn(nn_hdim, nn_output_dim) / np.sqrt(nn_hdim)).astype('float32'))\n",
    "    b2.set_value(np.zeros(nn_output_dim).astype('float32'))\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()\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())\n",
    "\n",
    "\n",
    "# Profiling\n",
    "# theano.config.profile = True\n",
    "# theano.config.profile_memory = True\n",
    "# gradient_step()\n",
    "# theano.printing.debugprint(gradient_step) \n",
    "# print gradient_step.profile.summary()\n",
    "            \n",
    "%timeit gradient_step()"
   ]
  }
 ],
 "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
}