notebooks/3-Instructor-deep-learning-from-scratch-pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deep learning from scratch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Learning objectives of the notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Get used to working with *PyTorch Tensors*, the core data structure needed for working with neural networks;\n",
    "- Practice using the `autograd` capabilities of PyTorch Tensors to carry out backpropagation without all the pain;\n",
    "- Apply the useful PyTorch `torch.no_grad` context manager for managing memory consumption;\n",
    "- Convert a NumPy-based gradient descent algorithm into one relying on PyTorch Tensors!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyTorch Basics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[PyTorch](http://pytorch.org) is a Python-based scientific computing package to support deep learning research. It provides tensor support (a replacement of NumPy, of sorts) to provide a fast & flexible platform for experimenting with neural networks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "print(f'PyTorch version: {torch.__version__}')\n",
    "print(f'NumPy version:   {np.__version__}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "The principal data structures in PyTorch are *tensors*; these are pretty much the same as standard multidimensional NumPy arrays. To illustrate this, let's construct a matrix of zeros (of `long` or 64 bit integer `dtype`) in NumPy, and then in PyTorch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# zeros construction in NumPy\n",
    "x_np = np.zeros((2,4), dtype=np.int64)\n",
    "print(x_np, x_np.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# zeros construction in PyTorch\n",
    "x = torch.zeros(2, 4, dtype=torch.long) # Observe difference in calling syntax!\n",
    "print(x, x.dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "You can query a tensor's size (dimensions) with the `size` method (contrast with NumPy array `shape` attribute)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "print(x)\n",
    "print(x.size())   # \"size\" is *method* for torch tensors\n",
    "print(x_np.shape) # 'shape' is *attribute* returning tuple\n",
    "print(x_np.size)  # \"size\"  is *attribute* for np arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# torch.Tensor.size() yields subclass of Python tuple\n",
    "print(type(x.size()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As with NumPy, there are a variety of PyTorch data types for arrays:\n",
    "\n",
    "|  NumPy dtype | PyTorch dtype | Alternative | Tensor class |\n",
    "|:-:|:-:|:-:|:-:|\n",
    "| `np.int16`  |`torch.Int16`  |`torch.short` |`ShortTensor` |\n",
    "| `np.int32`  |`torch.Int32`  |`torch.Int`   |`IntTensor`   |\n",
    "| `np.int64`  |`torch.Int64`  |`torch.long`  |`LongTensor`  |\n",
    "| `np.float16`|`torch.float16`|`torch.half`  |`HalfTensor`  |\n",
    "| `np.float32`|`torch.float32`|`torch.float` |`FloatTensor` |\n",
    "| `np.float64`|`torch.float64`|`torch.double`|`DoubleTensor`|\n",
    "\n",
    "\n",
    "Many functions and methods in PyTorch have similar names to NumPy functions & methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(torch.empty(3, 4, dtype=torch.short), end='\\n\\n')  # like numpy.empty\n",
    "print(torch.ones(3, 4, dtype=torch.short), end='\\n\\n')   # like numpy.ones\n",
    "print(torch.randn(3, 4, dtype=torch.float), end='\\n\\n')  # like numpy.random.randn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also construct PyTorch tensors from lists of numerical data or NumPy arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Constructing tensors from lists of data\n",
    "print(torch.tensor([1,2,3]).dtype)  # inferred to be 64 bit integers\n",
    "print(torch.Tensor([1,2,3]).dtype)  # specifically cast to 32 bit floats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice the factory function `torch.tensor` differs from the class constructor `torch.Tensor`. The former *infers* the data type of the tensor to construct from the numerical data input. By constrast, the latter is  just an alias for `Torch.FloatTensor` (i.e., the data are cast to 32 bit floating point numbers)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "PyTorch Tensors can be converted to NumPy arrays using the method `torch.Tensor.numpy`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "a = torch.rand(2,3)   # first, construct a random PyTorch tensor\n",
    "print(a)\n",
    "print(a.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "b = a.numpy()        # converts to NumPy array (shallow copy; use .copy() for deep copy)\n",
    "print(b)\n",
    "print(type(b))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[What is PyTorch?](https://pytorch.org/tutorials/beginner/blitz/tensor_tutorial.html)  at [`pytorch.org`](https://pytorch.org) provides a quick tour through related topics (e.g., tensor indexing, arithmetic operations, elementwise functions, linear algebra, etc.). For the most part, these resemble (although not perfectly) the same corresponding tasks in NumPy.\n",
    "\n",
    "# Backpropagation with `autograd`\n",
    "\n",
    "Why PyTorch Tensors when all they seem to offer is the same functionality of NumPy arrays? Another related question is why go through the trouble to reimplement everything that's done in NumPy in PyTorch (with slightly different names & APIs)? There are two principle advantages that systems like PyTorch have over NumPy for numerical computing:\n",
    "\n",
    "1. **Automatic differentiation**: PyTorch includes a package called [`autograd`](https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html) that computes the backpropagation algorithm for users. As such, the management of gradients (and the associated memory needed) is significantly simplified with the PyTorch framework. This is, of course, very important for implementing gradient descent.\n",
    "2. **GPU computation**: GPUs (graphical processing units) are widely available to speed up computation. However, GPU programming remains challenging for most developers with the memory management issues associated with moving data onto GPUs to speed up computation. With PyTorch, much of the work of moving tensors onto GPUs is handled for the user which makes programing with GPUs much easier... and this in turn speeds up a lot of neural network training.\n",
    "\n",
    "If we examine the object `a` created above, you can see it has an attribute `device` that can be set in various ways depending on the availability of GPU hardware."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a.device # PyTorch tensors have a `device` attribute\n",
    "# Common alternatives: device(type='cpu'), device(type='cuda'), etc."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "toc-hr-collapsed": true
   },
   "source": [
    "We'll focus mostly on automatic differentiation today as supported by the [`autograd`](https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html) module. Remember, our main reason for wanting to do this is to compute gradients as needed to train neural network parameters (weights & biases) with gradient descent. In PyTorch, automatic differentiation of tensors is achieved using through setting the `requires_grad` attribute to `True` for all relevant `torch.Tensor`s on construction (the default value is `False`). Alternatively, there is also a method `.requires_grad_( ... )` that modifies the `requires_grad` flag in-place (default value `False`).\n",
    "\n",
    "Once tensors are defined with the `requires_grad` attribute set correctly, additional space is allocated for intermediate computations (remember all the extra lists of arrays we had to maintain explicitly within the `forward` and `backward` functions?). These are used when calling `torch.Tensor.backward()` to compute all gradients recursively. The intermediate gradients computed can then be retrieved using the attribute `torch.Tensor.grad`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Backpropagation example\n",
    "\n",
    "Let's consider a simple polynomial function like below applied to a scalar value $x$:\n",
    "\n",
    "$\\begin{aligned} &\\mathrm{Function:} & f(x) &= 3x^4 -2x^3 + 4x^2 - x + 5 \\\\\n",
    "&\\mathrm{Derivative:} & f'(x) &= 12x^3 -6 x^2 + 8x -1\\end{aligned}$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "1. Create tensor `x` with the attribute `requires_grad=True` set in the constructor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "x = torch.tensor(2.0, requires_grad=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "2. Map the polynomial function $f$ onto tensor the `x` and assign the result to `y`. You can verify explicitly that, when $x=2$, $f(x)=51$:\n",
    " $$f(2)=3(2)^4 - 2(2)^3 + 4(2)^2 -(2) +5 = 48-16+16-2+5 = 51$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "y = 3*x**4 - 2*x**3 + 4*x**2 - x + 5  # Write out computation of y explicitly.\n",
    "\n",
    "print(y) # Notice y has a new attribute: grad_fn\n",
    "print(y.grad_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "The object `y` has an associated gradient function accessible as `y.grad_fn`. When `y` is computed and stored, a set of algebraic operations is applied to the tensor `x`. If the derivatives of those operations are known, the `autograd` package provides support for computing those derivatives (that's what the `AddBackward0` object is). Invoking `y.backward()`, then, computes the value of *gradient* of `y` with respect to `x` evaluated at `x==2`:\n",
    "\n",
    "$$f'(2) = 12(2^3) - 6(2^2) + 8(2) - 1 = 96-24+16-1 = 87. $$\n",
    "\n",
    "Notice that the computed gradient value is stored in the attribute `x.grad` of the original tensor `x`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 2,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "y.backward() # Compute derivatives and propagate values back through tensors on which y depends\n",
    "\n",
    "print(x.grad)  # Expect the value 87 as a singleton tensor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that invoking `y.backward()` a second time raises an exception. This is because the intermediate arrays required to execute the backpropagation have been released (i.e., the memory has been deallocated)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 2,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "y.backward() # Yields a RuntimeError (just like before calling backward before forward)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Another backpropagation example\n",
    "\n",
    "+ Use $z = \\cos(u)$ with $u=x^2$ at $x=\\sqrt{\\frac{\\pi}{3}}$\n",
    "+ Expect $z=\\frac{1}{2}$ when $x=\\sqrt{\\frac{\\pi}{3}}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "x = torch.tensor([np.sqrt(np.pi/3)], requires_grad=True)\n",
    "u = x ** 2\n",
    "z = torch.cos(u)\n",
    "print(f'x: {x}\\nu: {u}\\nz: {z}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "+ Expect \n",
    "  $$\\frac{dz}{dx} = \\frac{dz}{du} \\frac{du}{dx} = (-\\sin u) (2 x) = \\sqrt{\\pi}$$\n",
    "  when $x=\\sqrt{\\frac{\\pi}{3}}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Now apply backward for backpropagation of derivate values\n",
    "z.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "print(f'x.grad:\\t\\t\\t\\t\\t\\t{x.grad}')\n",
    "x, u = x.item(), u.item() # extract scalar values\n",
    "print(f'Computed derivative using analytic formula:\\t{-np.sin(u)*2*x}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that the tensors `x`, `u`, and `z` are all singleton tensors. The method `item` is used to extract a scalar entry out of a singleton tensor."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building a Neural Network in PyTorch\n",
    "\n",
    "Let's now use an approach adapted from one by [Justin Johnson](https://pytorch.org/tutorials/beginner/pytorch_with_examples.html)  (BSD Clause-3 License). The goal is to convert a NumPy-constructed gradient descent process modelling a feed-forward neural network into a PyTorch neural network. The architecture is similar to the one constructed in the last notebook.\n",
    "\n",
    "+ The input vectors are assumed to have $784(=28^2)$ features.\n",
    "+ The first layer is a hidden layer with 100 units and a *rectified linear unit* activation function (often called $\\mathrm{ReLU}$):\n",
    "\n",
    "$$ \\mathrm{ReLU}(x) = \\begin{cases} x, & \\mathrm{if\\ }x>0 \\\\ 0 & \\mathrm{otherwise} \\end{cases} \\quad\\Rightarrow\\quad\n",
    "\\mathrm{ReLU}'(x) = \\begin{cases} 1, & \\mathrm{if\\ }x>0 \\\\ 0 & \\mathrm{otherwise} \\end{cases}.\n",
    "$$\n",
    "\n",
    "+ The final output layer has 10 units and the activation function assiciated with this layer is the identity map.\n",
    "\n",
    "The loop provided below does not use functions to represent the initialization, forward propagation, backpropagation, and update steps of the steepest descent process. You'll use this as a starting point to develop a PyTorch version of this gradient descent loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_batch, dimensions = 64, [784, 100, 10]\n",
    "\n",
    "# Create random input and output data\n",
    "X = np.random.randn(dimensions[0], N_batch)\n",
    "y = np.random.randn(dimensions[-1], N_batch)\n",
    "\n",
    "# Randomly initialize weights & biases\n",
    "W1 = np.random.randn(dimensions[1], dimensions[0])\n",
    "W2 = np.random.randn(dimensions[2], dimensions[1])\n",
    "b1 = np.random.randn(dimensions[1], 1)\n",
    "b2 = np.random.randn(dimensions[2], 1)\n",
    "\n",
    "eta, MAXITER, SKIP = 5e-6, 2500, 100\n",
    "for epoch in range(MAXITER):\n",
    "    # Forward propagation: compute predicted y\n",
    "    Z1 = np.dot(W1, X) + b1\n",
    "    A1 = np.maximum(Z1, 0) # ReLU function\n",
    "    Z2 = np.dot(W2, A1) + b2\n",
    "    A2 = Z2   # identity function on output layer\n",
    "    \n",
    "    # Compute and print loss\n",
    "    loss = 0.5 * np.power(A2 - y, 2).sum()\n",
    "    if (divmod(epoch, SKIP)[1]==0):\n",
    "        print(epoch, loss)\n",
    "\n",
    "    # Backpropagation to compute gradients of loss with respect to W1, W2, b1, and b2\n",
    "    delta2 = (A2 - y)               # derivative of identity map == multiplying by ones\n",
    "    grad_W2 = np.dot(delta2, A1.T)\n",
    "    grad_b2 = np.dot(delta2, np.ones((N_batch, 1)))\n",
    "    delta1 = np.dot(W2.T, delta2) * (Z1>0) # derivative of ReLU is a step function\n",
    "    grad_W1 = np.dot(delta1, X.T)\n",
    "    grad_b1 = np.dot(delta1, np.ones((N_batch, 1)))\n",
    "\n",
    "    # Update weights & biases\n",
    "    W1 = W1 - eta * grad_W1\n",
    "    b1 = b1 - eta * grad_b1\n",
    "    W2 = W2 - eta * grad_W2\n",
    "    b2 = b2 - eta * grad_b2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Convert the preceding code to use PyTorch Tensors instead of NumPy arrays\n",
    "\n",
    "+ Replace use of `numpy.random.randn` with `torch.randn` to initialize the problem with PyTorch Tensors rather than NumPy arrays.\n",
    "+ Replace instances of `np.dot` with [`torch.mm`](https://pytorch.org/docs/stable/torch.html#torch.mm) (both of which implement standard matrix-vector products).\n",
    "+ Replace use of `np.ones` [`torch.ones`](https://pytorch.org/docs/stable/torch.html#torch.ones).\n",
    "+ Replace the computation of `A1 = np.maximum(Z1, 0)` with a call to the PyTorch builtin function `torch.relu`.\n",
    "+ Modify the computation of the `loss` to use PyTorch specific functions/methods (hint: there is a PyTorch `torch.Tensor.pow` method).\n",
    "+ When printing the loss every hundred epochs, use the `.item()` method to extract its singleton scalar entry.\n",
    "+ Make sure the loop executes in a similar fashion to the preceding loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_batch, dimensions = 64, [784, 100, 10]\n",
    "\n",
    "# Create random input and output data\n",
    "X = torch.randn(dimensions[0], N_batch)\n",
    "y = torch.randn(dimensions[-1], N_batch)\n",
    "\n",
    "# Randomly initialize weights & biases\n",
    "W1 = torch.randn(dimensions[1], dimensions[0])\n",
    "W2 = torch.randn(dimensions[2], dimensions[1])\n",
    "b1 = torch.randn(dimensions[1], 1)\n",
    "b2 = torch.randn(dimensions[2], 1)\n",
    "\n",
    "eta, MAXITER, SKIP = 5e-6, 2500, 100\n",
    "for epoch in range(MAXITER):\n",
    "    # Forward propagation: compute predicted y\n",
    "    Z1 = torch.mm(W1, X) + b1\n",
    "    A1 = torch.relu(Z1)        # Native PyTorch ReLU function\n",
    "    Z2 = torch.mm(W2, A1) + b2\n",
    "    A2 = Z2\n",
    "\n",
    "    # Compute and print loss\n",
    "    loss = 0.5 * (A2 - y).pow(2).sum()\n",
    "    if (divmod(epoch, SKIP)[1]==0):\n",
    "        print(epoch, loss.item())\n",
    "\n",
    "    # Backpropagation to compute gradients of loss with respect to W1, W2, b1, and b2\n",
    "    delta2 = (A2 - y)               # derivative of identity map == multiplying by ones\n",
    "    grad_W2 = torch.mm(delta2, A1.T)\n",
    "    grad_b2 = torch.mm(delta2, torch.ones(N_batch, 1))\n",
    "    delta1 = torch.mm(W2.T, delta2) * (Z1>0) # derivative of ReLU is a step function\n",
    "    grad_W1 = torch.mm(delta1, X.T)\n",
    "    grad_b1 = torch.mm(delta1, torch.ones(N_batch, 1))\n",
    "\n",
    "    # Update weights & biases\n",
    "    W1 = W1 - eta * grad_W1\n",
    "    b1 = b1 - eta * grad_b1\n",
    "    W2 = W2 - eta * grad_W2\n",
    "    b2 = b2 - eta * grad_b2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Use `backward()` and `grad` to compute backpropagation and updates\n",
    "\n",
    "Having set up the main loop with PyTorch Tensors, now make use of `autograd` to eliminate the tedious work of having to write the code to compute the gradients of the loss function with respect to `W1`, `W2`, `b1`, and `b2` explicitly.\n",
    "+ Insert `requires_grad=True` as an argument in the construction of `W1`, `W2`, `b1`, and `b2`.\n",
    "+ After computing the loss function value `loss`, replace all the lines used to compute gradients explicitly by a single call to `loss.backward()`.\n",
    "+ Replace the update steps with gradients stored in `.grad` attributes of the weights & biases. For instance, you can now compute `W1 -= eta * W1.grad` *after* the call to `loss.backward()` rather than computing and explicitly storing `grad_W1` and later computing `W1 -= eta * grad_W1`.\n",
    "    + Do these update steps within a `with torch.no_grad():` block (as provided below). The purpose of the [`torch.no_grad`](https://pytorch.org/docs/stable/torch.html#torch.no_grad) context manager is to reduce memory consumption.\n",
    "    + After completing the updates, zero out the computed gradients before the next iteration by calling the method `.zero_()`. For instance, you would call `W1.grad.zero_()` to zero out the computed gradient in place. This call will be within the scope of the `torch.no_grad` context manager.\n",
    "\n",
    "Notice, in PyTorch, methods like `.zero_` that have a training underscore in their name operate in place, i.e., they overwrite the memory locations associated with the tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create random input and output data\n",
    "X = torch.randn(dimensions[0], N_batch)\n",
    "y = torch.randn(dimensions[-1], N_batch)\n",
    "\n",
    "# Randomly initialize weights & biases\n",
    "W1 = torch.randn(dimensions[1], dimensions[0], requires_grad=True)\n",
    "W2 = torch.randn(dimensions[2], dimensions[1], requires_grad=True)\n",
    "b1 = torch.randn(dimensions[1], 1, requires_grad=True)\n",
    "b2 = torch.randn(dimensions[2], 1, requires_grad=True)\n",
    "\n",
    "eta, MAXITER, SKIP = 5e-6, 2500, 100\n",
    "for epoch in range(MAXITER):\n",
    "    # Forward propagation: compute predicted y\n",
    "    Z1 = torch.mm(W1, X) + b1\n",
    "    A1 = torch.relu(Z1)        # Native PyTorch ReLU function\n",
    "    Z2 = torch.mm(W2, A1) + b2\n",
    "    A2 = Z2\n",
    "\n",
    "    # Compute and print loss\n",
    "    loss = 0.5 * (A2 - y).pow(2).sum()\n",
    "    if (divmod(epoch, SKIP)[1]==0):\n",
    "        print(epoch, loss.item())\n",
    "\n",
    "    # Backpropagation to compute gradients of loss with respect to W1, W2, b1, and b2\n",
    "    loss.backward()\n",
    "\n",
    "    # Update weights & biases\n",
    "    with torch.no_grad():\n",
    "        W1 -= eta * W1.grad\n",
    "        b1 -= eta * b1.grad\n",
    "        W2 -= eta * W2.grad\n",
    "        b2 -= eta * b2.grad\n",
    "        # Manually zero the gradients after updating weights\n",
    "        W1.grad.zero_()\n",
    "        W2.grad.zero_()\n",
    "        b1.grad.zero_()\n",
    "        b2.grad.zero_()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# What next?\n",
    "\n",
    "PyTorch has a large ecosystem of utilities including packages like  `torch.nn` (which is like Keras in spirit to simplify specifying a network architecture in an object-oriented way) and `torch.optim` (which makes managing different optimization schemes easier). We've covered a lot of ground in this tutorial so far, so this will be as far as we can get today. But you now should have enough of an understanding of backpropagation that you can pick up more at [`pytorch.org`](https://pytorch.org) independently."
   ]
  }
 ],
 "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}