P14-C03-multiple-gpus-gluon.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#  Training with multiple GPUs with `gluon`\n",
    "\n",
    "Now we are going to implement the data parallelism algorithm introduced in [multi gpu from scratch](./P14-C02-multiple-gpus-scratch.ipynb) with `gluon`.\n",
    "\n",
    "First we define the example network and loss function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from mxnet import gluon, gpu\n",
    "net = gluon.nn.Sequential(prefix='cnn_')\n",
    "with net.name_scope():\n",
    "    net.add(gluon.nn.Conv2D(channels=20, kernel_size=3, activation='relu'))\n",
    "    net.add(gluon.nn.MaxPool2D(pool_size=(2,2), strides=(2,2)))\n",
    "    net.add(gluon.nn.Conv2D(channels=50, kernel_size=5, activation='relu'))\n",
    "    net.add(gluon.nn.MaxPool2D(pool_size=(2,2), strides=(2,2)))\n",
    "    net.add(gluon.nn.Flatten())\n",
    "    net.add(gluon.nn.Dense(128, activation=\"relu\"))\n",
    "    net.add(gluon.nn.Dense(10))\n",
    "    \n",
    "loss = gluon.loss.SoftmaxCrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize on multiple devices\n",
    "\n",
    "`gluon` supports initialize the network parameters over multiple devices. The parameters on one device is identical to the ones on another device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ctx = [gpu(0), gpu(1)]\n",
    "net.collect_params().initialize(ctx=ctx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given input data, the parameters on the according device are used to compute the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[[-0.01017658  0.03012515  0.02999702  0.01175333 -0.01746453  0.00707828\n",
      "   0.02404996  0.00616632 -0.02094562  0.0136827 ]\n",
      " [-0.01249129  0.0305641   0.02823936 -0.00159418 -0.00722831  0.00538148\n",
      "   0.01476716  0.0225275  -0.02458289  0.0246105 ]]\n",
      "<NDArray 2x10 @gpu(0)>\n",
      "\n",
      "[[-0.00349744  0.01896121  0.02959755  0.00261514  0.00015916 -0.00355723\n",
      "   0.0040103   0.03075583 -0.00761715  0.00599077]\n",
      " [-0.00557119  0.02766508  0.02406837 -0.0007478  -0.00511122  0.00538528\n",
      "   0.00292899  0.01488838 -0.00191687  0.01074106]]\n",
      "<NDArray 2x10 @gpu(1)>\n"
     ]
    }
   ],
   "source": [
    "from mxnet.test_utils import get_mnist\n",
    "mnist = get_mnist()\n",
    "batch = mnist['train_data'][0:4, :]\n",
    "data = gluon.utils.split_and_load(batch, ctx)\n",
    "print(net(data[0]))\n",
    "print(net(data[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can access the parameters on each device. (Note that, the weights may be initialized at the beginning of the first forward, while not in `initialize` because the data shapes may be not available at that time). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== channel 0 of the first conv2d on gpu(0) ===\n",
      "[[[ 0.0068339   0.01299825  0.0301265 ]\n",
      "  [ 0.04819721  0.01438687  0.05011239]\n",
      "  [ 0.00628365  0.04861524 -0.01068833]]]\n",
      "<NDArray 1x3x3 @gpu(0)>\n",
      "=== channel 0 of the first conv2d on gpu(1) ===\n",
      "[[[ 0.0068339   0.01299825  0.0301265 ]\n",
      "  [ 0.04819721  0.01438687  0.05011239]\n",
      "  [ 0.00628365  0.04861524 -0.01068833]]]\n",
      "<NDArray 1x3x3 @gpu(1)>\n"
     ]
    }
   ],
   "source": [
    "weight = net.collect_params()['cnn_conv2d0_weight']\n",
    "\n",
    "for c in ctx:\n",
    "    print('=== channel 0 of the first conv2d on {} ==={}'.format(\n",
    "        c, weight.data(ctx=c)[0]))\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similar we can access the gradients on each GPUs. Because the input data are different, the gradients on each GPU vary. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== grad of channel 0 of the first conv2d on gpu(0) ===\n",
      "[[[-0.00481181  0.02549154  0.05066926]\n",
      "  [ 0.01503928  0.04740802  0.04111018]\n",
      "  [ 0.04527877  0.06305876  0.04087965]]]\n",
      "<NDArray 1x3x3 @gpu(0)>\n",
      "=== grad of channel 0 of the first conv2d on gpu(1) ===\n",
      "[[[-0.01102538 -0.02251887 -0.02211753]\n",
      "  [-0.01587106 -0.03848278 -0.03960424]\n",
      "  [-0.03371563 -0.06092874 -0.064744  ]]]\n",
      "<NDArray 1x3x3 @gpu(1)>\n"
     ]
    }
   ],
   "source": [
    "def forward_backward(net, data, label):\n",
    "    with gluon.autograd.record():\n",
    "        losses = [loss(net(X), Y) for X, Y in zip(data, label)]\n",
    "    for l in losses:\n",
    "        l.backward()\n",
    "        \n",
    "label = gluon.utils.split_and_load(mnist['train_label'][0:4], ctx)\n",
    "forward_backward(net, data, label)\n",
    "for c in ctx:\n",
    "    print('=== grad of channel 0 of the first conv2d on {} ==={}'.format(\n",
    "        c, weight.grad(ctx=c)[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Put all things together\n",
    "\n",
    "Now we can implement the remaining functions. Most of them are the same as the previous tutorial, one notable difference is that a `gluon` trainer recognizes multi-devices, it will automatically aggregate the gradients and synchronize the parameters. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on [gpu(0)]\n",
      "Batch size is 64\n",
      "Epoch 0, training time = 4.9 sec\n",
      "         validation accuracy = 0.9764\n",
      "Epoch 1, training time = 4.7 sec\n",
      "         validation accuracy = 0.9839\n",
      "Epoch 2, training time = 4.6 sec\n",
      "         validation accuracy = 0.9836\n",
      "Epoch 3, training time = 4.7 sec\n",
      "         validation accuracy = 0.9861\n",
      "Epoch 4, training time = 4.7 sec\n",
      "         validation accuracy = 0.9862\n",
      "Running on [gpu(0), gpu(1)]\n",
      "Batch size is 128\n",
      "Epoch 0, training time = 2.8 sec\n",
      "         validation accuracy = 0.8951\n",
      "Epoch 1, training time = 2.8 sec\n",
      "         validation accuracy = 0.9687\n",
      "Epoch 2, training time = 2.8 sec\n",
      "         validation accuracy = 0.9759\n",
      "Epoch 3, training time = 2.8 sec\n",
      "         validation accuracy = 0.9785\n",
      "Epoch 4, training time = 2.9 sec\n",
      "         validation accuracy = 0.9813\n"
     ]
    }
   ],
   "source": [
    "from mxnet import nd\n",
    "from mxnet.io import NDArrayIter\n",
    "from time import time\n",
    "\n",
    "def train_batch(batch, ctx, net, trainer):\n",
    "    # split the data batch and load them on GPUs\n",
    "    data = gluon.utils.split_and_load(batch.data[0], ctx)\n",
    "    label = gluon.utils.split_and_load(batch.label[0], ctx)\n",
    "    # compute gradient\n",
    "    forward_backward(net, data, label)\n",
    "    # update parameters\n",
    "    trainer.step(batch.data[0].shape[0])\n",
    "    \n",
    "def valid_batch(batch, ctx, net):\n",
    "    data = batch.data[0].as_in_context(ctx[0])\n",
    "    pred = nd.argmax(net(data), axis=1)\n",
    "    return nd.sum(pred == batch.label[0].as_in_context(ctx[0])).asscalar()    \n",
    "\n",
    "def run(num_gpus, batch_size, lr):    \n",
    "    # the list of GPUs will be used\n",
    "    ctx = [gpu(i) for i in range(num_gpus)]\n",
    "    print('Running on {}'.format(ctx))\n",
    "    \n",
    "    # data iterator\n",
    "    mnist = get_mnist()\n",
    "    train_data = NDArrayIter(mnist[\"train_data\"], mnist[\"train_label\"], batch_size)\n",
    "    valid_data = NDArrayIter(mnist[\"test_data\"], mnist[\"test_label\"], batch_size)\n",
    "    print('Batch size is {}'.format(batch_size))\n",
    "    \n",
    "    net.collect_params().initialize(ctx=ctx)\n",
    "    trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': lr})\n",
    "    for epoch in range(5):\n",
    "        # train\n",
    "        start = time()\n",
    "        train_data.reset()\n",
    "        for batch in train_data:\n",
    "            train_batch(batch, ctx, net, trainer)\n",
    "        nd.waitall()  # wait all computations are finished to benchmark the time\n",
    "        print('Epoch %d, training time = %.1f sec'%(epoch, time()-start))\n",
    "        \n",
    "        # validating\n",
    "        valid_data.reset()\n",
    "        correct, num = 0.0, 0.0\n",
    "        for batch in valid_data:\n",
    "            correct += valid_batch(batch, ctx, net)\n",
    "            num += batch.data[0].shape[0]                \n",
    "        print('         validation accuracy = %.4f'%(correct/num))\n",
    "        \n",
    "run(1, 64, .3)        \n",
    "run(2, 128, .6)            "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "Both parameters and trainers in `gluon` supports multi-devices, moving from one device to multi-devices is straightforward. "
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}