P03.5-C02-custom-layer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Designing a custom layer with ``gluon`` \n",
    "\n",
    "Now that we've peeled back some of the syntactic sugar conferred by ``nn.Sequential()`` and given you a feeling for how ``gluon`` works under the hood, you might feel more comfortable when writing your high-level code. But the real reason to get to know ``gluon`` more intimately is so that we can mess around with it and write our own Blocks. \n",
    "\n",
    "Up until now, we've presented two versions of each tutorial. One from scratch and one in ``gluon``. Empowered with such independence, you might be wondering, \"if I wanted to write my own layer, why wouldn't I just do it from scratch?\" \n",
    "\n",
    "In reality, writing every model completely from scratch can be cumbersome.  Just like there's only so many times a developer can code up a blog from scratch without hating life, there's only so many times that you'll want to write out a convolutional layer, or define the stochastic gradient descent updates. Even in pure research environment, we usually want to customize one part of the model. For example, we might want to implement a new layer, but still rely on other common layers, loss functions, optimizers, etc. \n",
    "\n",
    "By hacking ``gluon``, we can get the desired flexibility in one part of our model, without screwing up everything else that makes our lives easy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import mxnet as mx\n",
    "from mxnet import nd, autograd, gluon\n",
    "from mxnet.gluon import nn, Block\n",
    "mx.random.seed(1)\n",
    "\n",
    "###########################\n",
    "#  Speficy the context we'll be using\n",
    "###########################\n",
    "ctx = mx.cpu()\n",
    "\n",
    "###########################\n",
    "#  Load up our dataset\n",
    "###########################\n",
    "batch_size = 64\n",
    "mnist = mx.test_utils.get_mnist()\n",
    "train_data = mx.io.NDArrayIter(\n",
    "    mnist[\"train_data\"], \n",
    "    mnist[\"train_label\"], \n",
    "    batch_size, \n",
    "    shuffle=True)\n",
    "test_data = mx.io.NDArrayIter(\n",
    "    mnist[\"test_data\"], \n",
    "    mnist[\"test_label\"], \n",
    "    batch_size, \n",
    "    shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining a (toy) custom layer\n",
    "\n",
    "To start, let's pretend that we want to use ``gluon`` for its optimizer, serialization, etc, but that we need a new layer. Specifically, we want a layer that centers its input about 0 by subtracting its mean. We'll go ahead and define the simplest possible ``Block``. Remember from the last tutorial that in ``gluon`` layers *are* ``Block``s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CenteredLayer(Block):\n",
    "    def __init__(self, **kwargs):\n",
    "        super(CenteredLayer, self).__init__(**kwargs)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return x - nd.mean(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it. We can just call instantiate this block and make a forward pass. \n",
    "Note that this layer doesn't actually care \n",
    "what it's input dimension or output dimensions are. \n",
    "So we can just feed in an arbitrary array \n",
    "and should expect appropriately transformed output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       "[-2. -1.  0.  1.  2.]\n",
       "<NDArray 5 @cpu(0)>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net = CenteredLayer()\n",
    "net(nd.array([1,2,3,4,5]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also incorporate this layer into a more complicated network, as by using ``nn.Sequential()``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "net2 = nn.Sequential()\n",
    "net2.add(nn.Dense(128))\n",
    "net2.add(nn.Dense(10))\n",
    "net2.add(CenteredLayer())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This network contains Blocks (Dense) that contain parameters and thus require initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "net2.collect_params().initialize(mx.init.Xavier(magnitude=2.24), ctx=ctx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can pass some data through it, say the first image from our MNIST dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[[ 0.59940612 -0.11329634 -0.94556594  0.15849063  0.59433293  0.12177399\n",
      "   0.00961533 -0.31887734  0.02881873 -0.13469785]]\n",
      "<NDArray 1x10 @cpu(0)>\n"
     ]
    }
   ],
   "source": [
    "output = net2(train_data.next().data[0][0])\n",
    "print(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can verify that as expected, the resulting vector has mean 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       "[  2.68220894e-08]\n",
       "<NDArray 1 @cpu(0)>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nd.mean(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There's a good chance you'll see something other than 0. When I ran this code, I got ``2.68220894e-08``. \n",
    "That's roughly ``.000000027``.\n",
    "This owes to the fact that MXNet gets performs low precision math. \n",
    "For deep learning research, this is often a compromise that we make.\n",
    "In exchange for giving up a few significant digits, we get tremendous speedups on modern hardware.\n",
    "And it turns out that most deep learning algorihtms don't work and worse on account of the loss of precision."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom layers with parameters\n",
    "\n",
    "While ``CenteredLayer`` should give you some sense of how to implement a custom layer, it's missing a few important pieces. Most importantly, ``CenteredLayer`` doesn't care about the dimensions of its input or output, and it doesn't contain any trainable parameters. Since you already know how to implement a fully-connected layer from scratch, let's learn how to make parameteric ``Block`` by implementing MyDense, our own version of a fully-connected (Dense) layer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameters\n",
    "\n",
    "Before we can add parameters to our custom ``Block``, we should get to know how ``gluon`` deals with parmaeters generally. Instead of working with NDArrays directly, each ``Block`` is associated with some number (as few as zero) of ``Parameter``s. \n",
    "\n",
    "At a high level, you can think of a ``Parameter`` as a wrapper on an ``NDArray``. However, the ``Parameter`` can be instantiated before the corresponding NDArray is. For example, when we instantiate a ``Block`` but the shapes of each parameter still need to be inferred, the Parameter will wait for the shape to be inferred before allocating memory. \n",
    "\n",
    "To get a hands-on feel for mxnet.Parameter, let's just instantiate one outside of a ``Block``:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter exciting_parameter_yay (shape=(5, 5), dtype=<class 'numpy.float32'>)\n"
     ]
    }
   ],
   "source": [
    "my_param = gluon.Parameter(\"exciting_parameter_yay\", grad_req='write', shape=(5,5))\n",
    "print(my_param)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we've instantiated a parameter, giving it the name \"exciting_parameter_yay\". We've also specified that we'll want to capture gradients for this Parameter. Under the hood, that let's ``gluon`` know that it has to call ``.attach_grad()`` on the underlying NDArray. We also specified the shape. Noew that we have a Parameter, we can initialize its values via ``.initialize()`` and print extract its data by calling ``.data()``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[[-0.15652239  0.25559878  0.61088544 -0.64541352  0.07847846]\n",
      " [-0.31043294 -0.1823152   0.50340134  0.4912681   0.27425849]\n",
      " [ 0.2392574   0.5435186  -0.52092296 -0.02812958 -0.62578702]\n",
      " [-0.13936114  0.29658896  0.28252971  0.16273445 -0.46218979]\n",
      " [-0.64452112  0.3151297  -0.42935488  0.35241729 -0.31812882]]\n",
      "<NDArray 5x5 @cpu(0)>\n"
     ]
    }
   ],
   "source": [
    "my_param.initialize(mx.init.Xavier(magnitude=2.24), ctx=ctx)\n",
    "print(my_param.data())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameter dictionaries (introducing ``ParameterDict``)\n",
    "\n",
    "Rather than directly store references to each of its ``Parameters``, ``Block``s typicaly contain a parameter dictionary (``ParameterDict``). In practice, we'll rarely instantiate our own ``ParamterDict``. That's because whenever we call the ``Block`` constructor it's generated automatically. For pedagogical purposes, we'll do it from scratch this one time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd = gluon.ParameterDict(prefix=\"block1_\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MXNet's ``ParameterDict`` does a few cool things for us. First, we can instantiate a new Parameter by calling ``pd.get()``"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter block1_exciting_parameter_yay (shape=(5, 5), dtype=<class 'numpy.float32'>)"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.get(\"exciting_parameter_yay\", grad_req='write', shape=(5,5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the new parameter is (i) contained in the ParameterDict and (ii) appends the prefix to it's name. This naming convention helps us to know which parameters belong to which ``Block`` or sub-``Block``. It's especially useful when we want to write parameters to disc, or read them from disc.\n",
    "\n",
    "Like a regular Python dictionary, we can get the names of all parameters with ``.keys()`` and can "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter block1_exciting_parameter_yay (shape=(5, 5), dtype=<class 'numpy.float32'>)"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd[\"block1_exciting_parameter_yay\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Craft a bespoke fully-connected ``gluon`` layer\n",
    "\n",
    "Now that we know how parameters work, we're ready to create our very own fully-connected layer. We'll use the familiar relu activation from previous tutorials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def relu(X):\n",
    "    return nd.maximum(X,nd.zeros_like(X))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can define our ``Block``. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyDense(Block):\n",
    "    ####################\n",
    "    # We add arguments to our constructor (__init__)\n",
    "    # to indicate the number of input units (``in_units``) \n",
    "    # and output units (``units``)\n",
    "    ####################\n",
    "    def __init__(self, units, in_units=0, **kwargs):\n",
    "        super(MyDense, self).__init__(**kwargs)\n",
    "        with self.name_scope():\n",
    "            self.units = units\n",
    "            self._in_units = in_units\n",
    "            #################\n",
    "            # We add the required parameters to the ``Block``'s ParameterDict , \n",
    "            # indicating the desired shape\n",
    "            #################\n",
    "            self.weight_initializer = mx.init.Xavier(magnitude=2.24)\n",
    "            self.weight = self.params.get(\n",
    "                'weight', init=self.weight_initializer, \n",
    "                shape=(in_units, units))\n",
    "            self.bias = self.params.get('bias', shape=(units,))        \n",
    "\n",
    "    #################\n",
    "    #  Now we just have to write the forward pass. \n",
    "    #  We could rely upong the FullyConnected primitative in NDArray, \n",
    "    #  but it's better to get our hands dirty and write it out\n",
    "    #  so you'll know how to compose arbitrary functions\n",
    "    #################\n",
    "    def forward(self, x):\n",
    "        with x.context:\n",
    "            linear = nd.dot(x, self.weight.data()) + self.bias.data()\n",
    "            activation = relu(linear)\n",
    "            return activation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall that every Block can be run just as if it were an entire network. \n",
    "In fact, linear models are nothing more than neural networks \n",
    "consisting of a single layer as a network.\n",
    "    \n",
    "So let's go ahead and run some data though our bespoke layer.\n",
    "We'll want to first instantiate the layer and initialize its parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "metadata": {},
   "outputs": [],
   "source": [
    "dense = MyDense(20, in_units=10)\n",
    "dense.collect_params().initialize(ctx=ctx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "mydense16_ (\n",
       "  Parameter mydense16_bias (shape=(20,), dtype=<class 'numpy.float32'>)\n",
       "  Parameter mydense16_weight (shape=(10, 20), dtype=<class 'numpy.float32'>)\n",
       ")"
      ]
     },
     "execution_count": 130,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dense.params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can run through some dummy data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       "[[ 0.          0.          0.41424903  0.14168519  0.39775169  1.14422572\n",
       "   1.35729218  0.2832939   0.68291545  0.52558887  0.          0.57638609\n",
       "   0.70538962  0.70836264  0.          0.          0.18341359  0.          0.\n",
       "   1.17828798]\n",
       " [ 0.          0.          0.41424903  0.14168519  0.39775169  1.14422572\n",
       "   1.35729218  0.2832939   0.68291545  0.52558887  0.          0.57638609\n",
       "   0.70538962  0.70836264  0.          0.          0.18341359  0.          0.\n",
       "   1.17828798]]\n",
       "<NDArray 2x20 @cpu(0)>"
      ]
     },
     "execution_count": 131,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dense(nd.ones(shape=(2,10)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using our layer to build an MLP\n",
    "\n",
    "While it's a good sanity check to run some data though the layer, the real proof that it works will be if we can compose a network entirely out of ``MyDense`` layers and achieve respectable accuracy on a real task. So we'll revist the MNIST digit classification task, and use the familiar ``nn.Sequential()`` syntax to build our net."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "net = gluon.nn.Sequential()\n",
    "with net.name_scope():\n",
    "    net.add(MyDense(128, in_units=784))\n",
    "    net.add(MyDense(64, in_units=128))\n",
    "    net.add(MyDense(10, in_units=64))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "metadata": {},
   "outputs": [],
   "source": [
    "net.collect_params().initialize(ctx=ctx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Instantiate a loss "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "loss = gluon.loss.SoftmaxCrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': .1})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation Metric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "metric = mx.metric.Accuracy()\n",
    "\n",
    "def evaluate_accuracy(data_iterator, net):\n",
    "    numerator = 0.\n",
    "    denominator = 0.\n",
    "    \n",
    "    data_iterator.reset()\n",
    "    for i, batch in enumerate(data_iterator):\n",
    "        with autograd.record():\n",
    "            data = batch.data[0].as_in_context(ctx).reshape((-1,784))\n",
    "            label = batch.label[0].as_in_context(ctx)\n",
    "            label_one_hot = nd.one_hot(label, 10)\n",
    "            output = net(data)\n",
    "        \n",
    "        metric.update([label], [output])\n",
    "    return metric.get()[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0. Loss: 0.453847481741, Train_acc 0.850856164384, Test_acc 0.852308917197\n",
      "Epoch 1. Loss: 0.371533283068, Train_acc 0.859275114155, Test_acc 0.852735623003\n",
      "Epoch 2. Loss: 0.333858854555, Train_acc 0.864212328767, Test_acc 0.860120632723\n",
      "Epoch 3. Loss: 0.312389601963, Train_acc 0.867562071918, Test_acc 0.864618136258\n",
      "Epoch 4. Loss: 0.297684676722, Train_acc 0.869900114155, Test_acc 0.867798793255\n",
      "Epoch 5. Loss: 0.0705092845548, Train_acc 0.887752092846, Test_acc 0.87265846946\n",
      "Epoch 6. Loss: 0.0565776970013, Train_acc 0.900970319635, Test_acc 0.889705199197\n",
      "Epoch 7. Loss: 0.0483166280769, Train_acc 0.9112175371, Test_acc 0.902392690488\n",
      "Epoch 8. Loss: 0.0414713247422, Train_acc 0.919414320142, Test_acc 0.91232512336\n",
      "Epoch 9. Loss: 0.0357186722606, Train_acc 0.92602739726, Test_acc 0.920256629545\n"
     ]
    }
   ],
   "source": [
    "epochs = 10\n",
    "moving_loss = 0.\n",
    "\n",
    "for e in range(epochs):\n",
    "    train_data.reset()\n",
    "    for i, batch in enumerate(train_data):\n",
    "        data = batch.data[0].as_in_context(ctx).reshape((-1,784))\n",
    "        label = batch.label[0].as_in_context(ctx)\n",
    "        with autograd.record():\n",
    "            output = net(data)\n",
    "            cross_entropy = loss(output, label)\n",
    "            cross_entropy.backward()\n",
    "        trainer.step(data.shape[0])\n",
    "        \n",
    "        ##########################\n",
    "        #  Keep a moving average of the losses\n",
    "        ##########################\n",
    "        if i == 0:\n",
    "            moving_loss = nd.mean(cross_entropy).asscalar()\n",
    "        else:\n",
    "            moving_loss = .99 * moving_loss + .01 * nd.mean(cross_entropy).asscalar()\n",
    "            \n",
    "    test_accuracy = evaluate_accuracy(test_data, net)\n",
    "    train_accuracy = evaluate_accuracy(train_data, net)\n",
    "    print(\"Epoch %s. Loss: %s, Train_acc %s, Test_acc %s\" % (e, moving_loss, train_accuracy, test_accuracy))    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "It works! There's a lot of other cool things you can do. In more advanced chapters, we'll show how you can make a layer that takes in multiple inputs, or one that cleverly calls down to MXNet's symbolic API to squeeze out extra peroformance without screwing up your convenient imperative workflow. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.4.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}