Introduction to PyTorch for former Torchies.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to PyTorch for former Torchies\n",
    "\n",
    "In this tutorial, you will learn the following:\n",
    "\n",
    "1. Using torch Tensors, and important difference against (Lua)Torch\n",
    "2. Using the autograd package\n",
    "3. Building neural networks\n",
    "  - Building a ConvNet\n",
    "  - Building a Recurrent Net\n",
    "  - Using multiple GPUs\n",
    "\n",
    "\n",
    "## Tensors \n",
    "\n",
    "Tensors behave almost exactly the same way in PyTorch as they do in Torch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 10\n",
      " 20\n",
      "[torch.LongStorage of size 2]\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "a = torch.FloatTensor(10, 20)\n",
    "# creates tensor of size (10 x 20) with uninitialized memory\n",
    "\n",
    "a = torch.randn(10, 20)\n",
    "# initializes a tensor randomized with a normal distribution with mean=0, var=1\n",
    "\n",
    "print(a.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inplace / Out-of-place\n",
    "\n",
    "The first difference is that ALL operations on the tensor that operate in-place on it will have an `_` postfix.\n",
    "For example, `add` is the out-of-place version, and `add_` is the in-place version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "a.fill_(3.5)\n",
    "# a has now been filled with the value 3.5\n",
    "\n",
    "b = a.add(4.0)\n",
    "# a is still filled with 3.5\n",
    "# new tensor b is returned with values 3.5 + 4.0 = 7.5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some operations like narrow do not have in-place versions, and hence, .narrow_ does not exist. \n",
    "Similarly, some operations like fill_ do not have an out-of-place version, so .fill does not exist.\n",
    "\n",
    "### Zero Indexing\n",
    "\n",
    "Another difference is that Tensors are zero-indexed. (Torch tensors are one-indexed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "b = a[0,3] # select 1st row, 4th column from a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tensors can be also indexed with Python's slicing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "b = a[:,3:5] # selects all rows, columns 3 to 5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### No camel casing\n",
    "\n",
    "The next small difference is that all functions are now NOT camelCase anymore.\n",
    "For example `indexAdd` is now called `index_add_`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " 1  1  1  1  1\n",
      " 1  1  1  1  1\n",
      " 1  1  1  1  1\n",
      " 1  1  1  1  1\n",
      " 1  1  1  1  1\n",
      "[torch.FloatTensor of size 5x5]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "x = torch.ones(5, 5)\n",
    "print(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "  10  100\n",
      "  10  100\n",
      "  10  100\n",
      "  10  100\n",
      "  10  100\n",
      "[torch.FloatTensor of size 5x2]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "z = torch.Tensor(5, 2)\n",
    "z[:,0] = 10\n",
    "z[:,1] = 100\n",
    "print(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " 101    1    1    1   11\n",
      " 101    1    1    1   11\n",
      " 101    1    1    1   11\n",
      " 101    1    1    1   11\n",
      " 101    1    1    1   11\n",
      "[torch.FloatTensor of size 5x5]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "x.index_add_(1, torch.LongTensor([4,0]), z)\n",
    "print(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Numpy Bridge\n",
    "\n",
    "Converting a torch Tensor to a numpy array and vice versa is a breeze.\n",
    "The torch Tensor and numpy array will share their underlying memory, and changing one will change the other.\n",
    "\n",
    "#### Converting torch Tensor to numpy Array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " 1\n",
      " 1\n",
      " 1\n",
      " 1\n",
      " 1\n",
      "[torch.FloatTensor of size 5]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "a = torch.ones(5)\n",
    "print(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 1.  1.  1.  1.  1.]\n"
     ]
    }
   ],
   "source": [
    "b = a.numpy()\n",
    "print(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " 2\n",
      " 2\n",
      " 2\n",
      " 2\n",
      " 2\n",
      "[torch.FloatTensor of size 5]\n",
      "\n",
      "[ 2.  2.  2.  2.  2.]\n"
     ]
    }
   ],
   "source": [
    "a.add_(1)\n",
    "print(a)\n",
    "print(b) # see how the numpy array changed in value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Converting numpy Array to torch Tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 2.  2.  2.  2.  2.]\n",
      "\n",
      " 2\n",
      " 2\n",
      " 2\n",
      " 2\n",
      " 2\n",
      "[torch.DoubleTensor of size 5]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "a = np.ones(5)\n",
    "b = torch.DoubleTensor(a)\n",
    "np.add(a, 1, out=a)\n",
    "print(a)\n",
    "print(b) # see how changing the np array changed the torch Tensor automatically"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All the Tensors on the CPU except a CharTensor support converting to NumPy and back.\n",
    "\n",
    "### CUDA Tensors\n",
    "\n",
    "CUDA Tensors are nice and easy in pytorch, and they are much more consistent as well.\n",
    "Transfering a CUDA tensor from the CPU to GPU will retain it's type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# creates a LongTensor and transfers it \n",
    "# to GPU as torch.cuda.LongTensor\n",
    "a = torch.LongTensor(10).fill_(3).cuda()\n",
    "print(type(a))\n",
    "b = a.cpu()\n",
    "# transfers it to CPU, back to \n",
    "# being a torch.LongTensor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Autograd\n",
    "\n",
    "Autograd is now a core torch package for automatic differentiation. \n",
    "\n",
    "It uses a tape based system for automatic differentiation. \n",
    "\n",
    "In the forward phase, the autograd tape will remember all the operations it executed, and in the backward phase, it will replay the operations.\n",
    "\n",
    "In autograd, we introduce a `Variable` class, which is a very thin wrapper around a `Tensor`. \n",
    "You can access the raw tensor through the `.data` attribute, and after computing the backward pass, a gradient w.r.t. this variable is accumulated into `.grad` attribute.\n",
    "\n",
    "![Variable](images/Variable.png)\n",
    "\n",
    "There's one more class which is very important for autograd implementation - a `Function`. `Variable` and `Function` are interconnected and build up an acyclic graph, that encodes a complete history of computation. Each variable has a `.creator` attribute that references a function that has created a function (except for Variables created by the user - these have `None` as  `.creator`).\n",
    "\n",
    "If you want to compute the derivatives, you can call `.backward()` on a `Variable`. \n",
    "If `Variable` is a scalar (i.e. it holds a one element tensor), you don't need to specify any arguments to `backward()`, however if it has more elements, you need to specify a `grad_output` argument that is a tensor of matching shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Variable containing:\n",
       " 1  1\n",
       " 1  1\n",
       "[torch.FloatTensor of size 2x2]"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torch.autograd import Variable\n",
    "x = Variable(torch.ones(2, 2))\n",
    "x  # notice the \"Variable containing\" line"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       " 1  1\n",
       " 1  1\n",
       "[torch.FloatTensor of size 2x2]"
      ]
     },
     "execution_count": 95,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       " 0  0\n",
       " 0  0\n",
       "[torch.FloatTensor of size 2x2]"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 97,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.creator is None  # we've created x ourselves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Variable containing:\n",
       " 3  3\n",
       " 3  3\n",
       "[torch.FloatTensor of size 2x2]"
      ]
     },
     "execution_count": 98,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y = x + 2\n",
    "y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch.autograd.functions.basic_ops.AddConstant at 0x10d4076d0>"
      ]
     },
     "execution_count": 99,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y.creator\n",
    "# y was created as a result of an operation, \n",
    "# so it has a creator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Variable containing:\n",
       " 27  27\n",
       " 27  27\n",
       "[torch.FloatTensor of size 2x2]"
      ]
     },
     "execution_count": 100,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z = y * y * 3\n",
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Variable containing:\n",
       " 27\n",
       "[torch.FloatTensor of size 1]"
      ]
     },
     "execution_count": 101,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out = z.mean()\n",
    "out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# let's backprop now\n",
    "out.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       " 4.5000  4.5000\n",
       " 4.5000  4.5000\n",
       "[torch.FloatTensor of size 2x2]"
      ]
     },
     "execution_count": 103,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# print gradients d(out)/dx\n",
    "x.grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, gradient computation flushes all the internal buffers contained in the graph, so if you even want to do the backward on some part of the graph twice, you need to pass in `retain_variables = True` during the first pass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       " 1  1\n",
       " 1  1\n",
       "[torch.FloatTensor of size 2x2]"
      ]
     },
     "execution_count": 104,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = Variable(torch.ones(2, 2))\n",
    "y = x + 2\n",
    "y.backward(torch.ones(2, 2), retain_variables=True)\n",
    "# the retain_variables flag will prevent the internal buffers from being freed\n",
    "x.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Variable containing:\n",
       " 9  9\n",
       " 9  9\n",
       "[torch.FloatTensor of size 2x2]"
      ]
     },
     "execution_count": 105,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z = y * y\n",
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       " 5.8201  9.2361\n",
       " 1.5865 -6.4018\n",
       "[torch.FloatTensor of size 2x2]"
      ]
     },
     "execution_count": 106,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gradient = torch.randn(2, 2)\n",
    "# just backproping random gradients\n",
    "\n",
    "z.backward(gradient)\n",
    "\n",
    "# this would fail if we didn't specify \n",
    "# that we want to retain variables\n",
    "x.grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## nn package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We've redesigned the nn package, so that it's much more intuitive and fully integrated with autograd.\n",
    "\n",
    "### Replace containers with autograd\n",
    "\n",
    "You no longer have to use Containers like ConcatTable, or modules like CAddTable, or use and debug with nngraph. \n",
    "We will seamlessly use autograd to define our neural networks.\n",
    "For example, \n",
    "\n",
    "`output = CAddTable():forward({input1, input2})` simply becomes `output = input1 + input2`\n",
    "\n",
    "`output = MulConstant(0.5):forward(input)` simply becomes `output = input * 0.5`\n",
    "\n",
    "### State is no longer held in the module, but in the network graph\n",
    "\n",
    "Using recurrent networks is a breeze. If you want to create a recurrent network, simply use the same Linear layer multiple times, and the weights are shared.\n",
    "\n",
    "![torch-nn-vs-pytorch-nn](images/torch-nn-vs-pytorch-nn.png)\n",
    "\n",
    "### Simplified debugging\n",
    "\n",
    "Debugging is intuitive using Python's pdb debugger, and **the debugger and stack traces stop at exactly where an error occurred.** What you see is what you get.\n",
    "\n",
    "### Example 1: ConvNet\n",
    "\n",
    "Creating networks is simple. All of your networks are derived from the base class nn.Container.\n",
    "\n",
    "- In the constructor, you declare all the layers you want to use.\n",
    "- In the forward function, you define how your model is going to be run, from input to output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class MNISTConvNet(nn.Container):\n",
    "    def __init__(self):\n",
    "        # this is the place where you instantiate all your modules\n",
    "        # you can later access them using the same names you've given them in here\n",
    "        super(MNISTConvNet, self).__init__(\n",
    "            conv1 = nn.Conv2d(1, 20, 5),\n",
    "            pool1 = nn.MaxPool2d(2, 2),\n",
    "            conv2 = nn.Conv2d(20, 50, 5),\n",
    "            pool2 = nn.MaxPool2d(2, 2),\n",
    "            fc1   = nn.Linear(800, 500),\n",
    "            fc2   = nn.Linear(500, 10),\n",
    "            relu  = nn.ReLU(),\n",
    "            softmax = nn.LogSoftmax(),\n",
    "        )\n",
    "        \n",
    "    # it's the forward function that defines the network structure\n",
    "    # we're accepting only a single input in here, but if you want,\n",
    "    # feel free to use more\n",
    "    def forward(self, input):\n",
    "        x = self.pool1(self.relu(self.conv1(input)))\n",
    "        x = self.pool2(self.relu(self.conv2(x)))\n",
    "\n",
    "        # in your model definition you can go full crazy and use arbitrary\n",
    "        # python code to define your model structure\n",
    "        # all these are perfectly legal, and will be handled correctly \n",
    "        # by autograd:\n",
    "        # if x.gt(0) > x.numel() / 2:\n",
    "        #      ...\n",
    "        # \n",
    "        # you can even do a loop and reuse the same module inside it\n",
    "        # modules no longer hold ephemeral state, so you can use them\n",
    "        # multiple times during your forward pass\n",
    "        # e.g. in this example we're using relu multiple times (but\n",
    "        # modules with parameters will have correct gradients w.r.t.\n",
    "        # their weights as well)\n",
    "        # while x.norm(2) < 10:\n",
    "        #    x = self.conv1(x) \n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.relu(self.fc1(x))\n",
    "        x = self.relu(self.fc2(x))\n",
    "        return self.softmax(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using this ConvNet now is intuitive. \n",
    "You create an instance of it, and you can run things through it and backward."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 1\n",
      " 10\n",
      "[torch.LongStorage of size 2]\n"
     ]
    }
   ],
   "source": [
    "net = MNISTConvNet()\n",
    "# only mini-batches are supported in all of nn\n",
    "input = Variable(torch.randn(1, 1, 28, 28))\n",
    "out = net(input)\n",
    "print(out.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Variable containing:\n",
      " 2.3291\n",
      "[torch.FloatTensor of size 1]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "target = Variable(torch.LongTensor([3]), requires_grad = False)\n",
    "loss_fn = nn.NLLLoss()\n",
    "err = loss_fn(out, target)\n",
    "print(err)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "err.backward()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's easy to access individual layer gradients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 20\n",
      " 1\n",
      " 5\n",
      " 5\n",
      "[torch.LongStorage of size 4]\n"
     ]
    }
   ],
   "source": [
    "print(net.conv1.weight.grad.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.451569274352\n"
     ]
    }
   ],
   "source": [
    "print(net.conv1.weight.grad.norm())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A full and working MNIST example is located here\n",
    "https://github.com/pytorch/examples/tree/master/mnist\n",
    "\n",
    "### Example 2: Recurrent Net\n",
    "\n",
    "Building recurrent nets with PyTorch is quite a breeze.\n",
    "Since the state of the network is held in the graph and not\n",
    "in the layers, you can simply create an nn.Linear and \n",
    "reuse it over and over again for the recurrence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class RNN(nn.Container):\n",
    "\n",
    "    # you can also accept arguments in your model constructor\n",
    "    # if you want to parametrize it somehow\n",
    "    def __init__(self, data_size, hidden_size, output_size):\n",
    "        self.hidden_size = hidden_size\n",
    "        input_size = data_size + hidden_size\n",
    "        super(RNN, self).__init__(\n",
    "            i2h=nn.Linear(input_size, hidden_size),\n",
    "            i2o=nn.Linear(input_size, output_size)\n",
    "        )\n",
    "    \n",
    "    def forward(self, data, last_hidden):\n",
    "        input = data.cat((data, last_hidden,), 1)\n",
    "        hidden = self.i2h(input)\n",
    "        output = self.i2o(input)\n",
    "        return hidden, output\n",
    "\n",
    "rnn = RNN(50, 20, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "loss_fn = nn.MSELoss()\n",
    "\n",
    "batch_size = 10\n",
    "TIMESTEPS = 5\n",
    "batch = Variable(torch.randn(batch_size, 50), requires_grad = False)\n",
    "hidden = Variable(torch.zeros(batch_size, 20), requires_grad = False)\n",
    "target = Variable(torch.zeros(batch_size, 10), requires_grad = False)\n",
    "loss = 0\n",
    "for t in range(TIMESTEPS):                  \n",
    "    # yes! you can reuse the same network several times,\n",
    "    # sum up the losses, and call backward!\n",
    "    hidden, output = rnn(batch, hidden)\n",
    "    loss += loss_fn(output, target)\n",
    "loss.backward()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-GPU examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Data Parallel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class DataParallelModel(nn.Container):\n",
    "    def __init__(self):\n",
    "        super().__init__(\n",
    "            block1=nn.Linear(10, 20),\n",
    "            block2=nn.Linear(20, 20),\n",
    "            block3=nn.Linear(20, 20),\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.block1(x)\n",
    "        x = nn.parallel.data_parallel(self.block2, x, (0, 1, 2, 3))\n",
    "        x = self.block3(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Model Parallel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class DistributedModel(nn.Container):\n",
    "    def __init__(self):\n",
    "        super().__init__(\n",
    "            embedding=nn.Embedding(1000, 10),\n",
    "            rnn=nn.Linear(10, 10).cuda(0),\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.embedding(x)\n",
    "        x = x.cuda(0)\n",
    "        x = self.rnn(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Primitives on which data parallel is implemented upon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def data_parallel(module, input, device_ids, output_device=None):\n",
    "    \"\"\"Distributes replicas of module accross gpus given in device_ids,\n",
    "       slices the input and applies the copies in parallel.\n",
    "       Outputs are concatenated on the same device as input (or on\n",
    "       output_device if specified). Device id -1 means the CPU.\n",
    "    \"\"\"\n",
    "    if not device_ids:\n",
    "        return module(input)\n",
    "    replicas = replicate(module, device_ids)\n",
    "    inputs = scatter(input, device_ids)\n",
    "    outputs = parallel_apply(replicas, inputs)\n",
    "    if output_device is None:\n",
    "        output_device = -1 if not input.is_cuda else input.get_device()\n",
    "    return gather(outputs, output_device)"
   ]
  },
  {
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}