Part1_Backprop.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\newcommand{\\mat}[1]{\\boldsymbol {#1}}\n",
    "\\newcommand{\\mattr}[1]{\\boldsymbol {#1}^\\top}\n",
    "\\newcommand{\\matinv}[1]{\\boldsymbol {#1}^{-1}}\n",
    "\\newcommand{\\vec}[1]{\\boldsymbol {#1}}\n",
    "\\newcommand{\\vectr}[1]{\\boldsymbol {#1}^\\top}\n",
    "\\newcommand{\\rvar}[1]{\\mathrm {#1}}\n",
    "\\newcommand{\\rvec}[1]{\\boldsymbol{\\mathrm{#1}}}\n",
    "\\newcommand{\\diag}{\\mathop{\\mathrm {diag}}}\n",
    "\\newcommand{\\set}[1]{\\mathbb {#1}}\n",
    "\\newcommand{\\norm}[1]{\\left\\lVert#1\\right\\rVert}\n",
    "\\newcommand{\\pderiv}[2]{\\frac{\\partial #1}{\\partial #2}}\n",
    "\\newcommand{\\bb}[1]{\\boldsymbol{#1}}\n",
    "$$\n",
    "# Part 1: Backpropagation\n",
    "<a id=part1></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this part we will learn about **backpropagation** and **automatic differentiation**. We'll implement both of these concepts from scratch and compare our implementation to `PyTorch`'s built in implementation (`autograd`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import unittest\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "test = unittest.TestCase()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The backpropagation algorithm is at the core of training deep models. To state the problem we'll tackle in this notebook, imagine we have an L-layer MLP model, defined as\n",
    "$$\n",
    "\\hat{\\vec{y}^i} = \\vec{y}_L^i= \\varphi_L \\left(\n",
    "\\mat{W}_L \\varphi_{L-1} \\left( \\cdots\n",
    "\\varphi_1 \\left( \\mat{W}_1 \\vec{x}^i + \\vec{b}_1 \\right)\n",
    "\\cdots \\right)\n",
    "+ \\vec{b}_L \\right),\n",
    "$$\n",
    "\n",
    "a pointwise loss function $\\ell(\\vec{y}, \\hat{\\vec{y}})$ and an empirical loss over our entire data set,\n",
    "$$\n",
    "L(\\vec{\\theta}) = \\frac{1}{N} \\sum_{i=1}^{N} \\ell(\\vec{y}^i, \\hat{\\vec{y}^i}) + R(\\vec{\\theta})\n",
    "$$\n",
    "\n",
    "where $\\vec{\\theta}$ is a vector containing all network parameters, e.g.\n",
    "$\\vec{\\theta} = \\left[ \\mat{W}_{1,:}, \\vec{b}_1, \\dots,  \\mat{W}_{L,:}, \\vec{b}_L \\right]$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to train our model we would like to calculate the derivative\n",
    "(or **gradient**, in the multivariate case) of the loss with respect to each and every one of the parameters,\n",
    "i.e. $\\pderiv{L}{\\mat{W}_j}$ and $\\pderiv{L}{\\vec{b}_j}$ for all $j$.\n",
    "Since the gradient \"points\" to the direction of functional increase, the negative gradient is often used as a descent direction for descent-based optimization algorithms.\n",
    "In other words, iteratively updating each parameter proportianally to it's negetive gradient can lead to\n",
    "convergence to a local minimum of the loss function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculus tells us that as long as we know the derivatives of all the functions \"along the way\"\n",
    "($\\varphi_i(\\cdot),\\ \\ell(\\cdot,\\cdot),\\ R(\\cdot)$)\n",
    "we can use the **chain rule** to calculate the derivative \n",
    "of the loss with respect to any one of the parameter vectors.\n",
    "Note that if the loss $L(\\vec{\\theta})$ is scalar (which is usually the case), the gradient of a parameter\n",
    "will have the same shape as the parameter itself (matrix/vector/tensor of same dimensions).\n",
    "\n",
    "For deep models that are a composition of many functions, calculating the gradient of each parameter by hand and implementing hard-coded gradient derivations quickly becomes infeasible.\n",
    "Additionally, such code makes models hard to change, since any change potentially requires re-derivation and re-implementation of the entire gradient function.\n",
    "\n",
    "The backpropagation algorithm, which we saw [in the lecture](https://vistalab-technion.github.io/cs236605/lecture_notes/lecture_3/#error-backpropagation), provides us with a effective method of applying the **chain rule** recursively so that we can implement gradient calculations of arbitrarily deep or complex models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll now implement backpropagation using a modular software design which will allow us to chain many components layers together and get automatic gradient calculation of the output with respect to the input or any intermediate parameter."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To do this, we'll define a (very poorly named) class: `Block`. Here's the API of this class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Help on class Block in module hw2.blocks:\n",
      "\n",
      "class Block(abc.ABC)\n",
      " |  A block is some computation element in a network architecture which\n",
      " |  supports automatic differentiation using forward and backward functions.\n",
      " |  \n",
      " |  Method resolution order:\n",
      " |      Block\n",
      " |      abc.ABC\n",
      " |      builtins.object\n",
      " |  \n",
      " |  Methods defined here:\n",
      " |  \n",
      " |  __call__(self, *args, **kwargs)\n",
      " |      Call self as a function.\n",
      " |  \n",
      " |  __init__(self)\n",
      " |      Initialize self.  See help(type(self)) for accurate signature.\n",
      " |  \n",
      " |  backward(self, dout)\n",
      " |      Computes the backward pass of the block, i.e. the gradient\n",
      " |      calculation of the final network output with respect to each of the\n",
      " |      parameters of the forward function.\n",
      " |      :param dout: The gradient of the network with respect to the\n",
      " |      output of this block.\n",
      " |      :return: A tuple with the same number of elements as the parameters of\n",
      " |      the forward function. Each element will be the gradient of the\n",
      " |      network output with respect to that parameter.\n",
      " |  \n",
      " |  forward(self, *args, **kwargs)\n",
      " |      Computes the forward pass of the block.\n",
      " |      :param args: The computation arguments (implementation specific).\n",
      " |      :return: The result of the computation.\n",
      " |  \n",
      " |  params(self)\n",
      " |      :return: Block's trainable parameters and their gradients as a list\n",
      " |      of tuples, each tuple containing a tensor and it's corresponding\n",
      " |      gradient tensor.\n",
      " |  \n",
      " |  train(self, training_mode=True)\n",
      " |      Changes the mode of this block between training and evaluation (test)\n",
      " |      mode. Some blocks have different behaviour depending on mode.\n",
      " |      :param training_mode: True: set the model in training mode. False: set\n",
      " |      evaluation mode.\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Data descriptors defined here:\n",
      " |  \n",
      " |  __dict__\n",
      " |      dictionary for instance variables (if defined)\n",
      " |  \n",
      " |  __weakref__\n",
      " |      list of weak references to the object (if defined)\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Data and other attributes defined here:\n",
      " |  \n",
      " |  __abstractmethods__ = frozenset({'backward', 'forward', 'params'})\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import hw2.blocks as blocks\n",
    "help(blocks.Block)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In other words, a `Block` can be anything: a layer, an activation function, a loss function or generally *any computation that we know how to derive a gradient for*.\n",
    "\n",
    "Each block must define a `forward()` function and a `backward()` function.\n",
    "- The `forward()` function performs the actual calculation/operation of the block and returns an output.\n",
    "- The `backward()` function computes the gradient of the input and parameters as a function of the gradient of the **output**, according to the chain rule.\n",
    "\n",
    "This may sound confusing, so here's a diagram illustrating what a `Block` does:\n",
    "\n",
    "<img src=\"imgs/backprop.png\" width=\"900\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the diagram doesn't show that if $f(\\vec{x},\\vec{y})$ is parametrized, there are also gradients to calculate for the parameters.\n",
    "\n",
    "The forward pass is pretty straightforward: just do the computation.\n",
    "To understand the backward pass, imagine that there's some \"downstream\" loss function\n",
    "$L(\\vec{\\theta})$ and magically somehow we are told the gradient of that loss with respect\n",
    "to the output of our block, $\\pderiv{L}{\\vec{z}}$.\n",
    "\n",
    "Now, since $f(\\vec{x},\\vec{y})$ is a function we know how to derivate,\n",
    "it means we know how to calculate $\\pderiv{\\vec{z}}{\\vec{x}}$, $\\pderiv{\\vec{z}}{\\vec{y}}$ and also the derivatives with respect to any parameters that $f$ uses.\n",
    "Thanks to the chain rule, this is all we need to calculate the gradients of the **loss** w.r.t. the input and\n",
    "parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparison with PyTorch\n",
    "<a id=part1_1></a>\n",
    "\n",
    "PyTorch has the `nn.Module` base class, which is similar to our `Block` since it also represents a computation element in a network.\n",
    "However PyTorch's `nn.Module`s don't compute the gradient,\n",
    "only perform the forward calculations.\n",
    "Instead, the `autograd` module tracks operations on tensors, creating a graph of operations, each with it's own `backward` function. Therefore, in PyTorch the `backward()` function is called on the tensors, not the modules.\n",
    "\n",
    "Here we'll implement a \"poor man's autograd\": We'll use PyTorch tensors,\n",
    "but we'll calculate and store the gradients in our blocks (or return them). The gradients we'll calculate are of the entire block, not individual operations on tensors.\n",
    "\n",
    "To test our  implementation, we'll use PyTorch's `autograd`.\n",
    "\n",
    "Note that of course this method of tracking gradients is **much** more limited than what PyTorch offers. However it allows us to implement the backpropagation algorithm very simply and really see how it works.\n",
    "\n",
    "Let's set up some testing instrumentation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mSignature:\u001b[0m \u001b[0mcompare_block_to_torch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mblock\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mhw2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mblocks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mBlock\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mseed\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m42\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mSource:\u001b[0m   \n",
       "\u001b[0;32mdef\u001b[0m \u001b[0mcompare_block_to_torch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mblock\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mblocks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mBlock\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mseed\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m42\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0;34m\"\"\"\u001b[0m\n",
       "\u001b[0;34m    Compares the manually calculated gradients of a Block (it's backward\u001b[0m\n",
       "\u001b[0;34m    function) to the gradients produced by PyTorch's autograd.\u001b[0m\n",
       "\u001b[0;34m    \"\"\"\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0;31m# Forward pass\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmanual_seed\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mseed\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mz\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mblock\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0;31m# Invent some output gradient\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mdz\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mz\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mz\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m0\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1.0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0;31m# Backward pass (ours)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mdx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mblock\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdz\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0;31m# Attach autograd gradients to params and input and re-run forward pass on\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0;31m# the same input\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0;32mfor\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m \u001b[0;32min\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mblock\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrequires_grad\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmanual_seed\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mseed\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mz\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mblock\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0;31m# Backward pass (this time with PyTorch autograd)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mz\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdz\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'Comparing gradients... '\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mdiffs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0;31m# Compare input gradient\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mdx_autograd\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgrad\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mdiffs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnorm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdx_autograd\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mdx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf'{\"input\":8s} diff={diffs[-1]:.3f}'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0;31m# Compare parameter gradients\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mp\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdp\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mblock\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mdp_autograd\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgrad\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mdiffs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnorm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdp_autograd\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mdp\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf'param#{i+1:02d} diff={diffs[-1]:.3f}'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0;32mreturn\u001b[0m \u001b[0mdiffs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mFile:\u001b[0m      ~/236605/cs236605-hw2/hw2/grad_compare.py\n",
       "\u001b[0;31mType:\u001b[0m      function\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from hw2.grad_compare import compare_block_to_torch\n",
    "\n",
    "def test_block_grad(block: blocks.Block, x, y=None, delta=1e-2):\n",
    "    diffs = compare_block_to_torch(block, x, y)\n",
    "    \n",
    "    # Assert diff values\n",
    "    for diff in diffs:\n",
    "        test.assertLess(diff, delta)\n",
    "\n",
    "# Show the compare function\n",
    "compare_block_to_torch??"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notes:\n",
    "- After you complete your implementation, you should make sure to read and understand the `compare_block_to_torch()` function. It will help you understand what PyTorch is doing.\n",
    "- The value of `delta` above is somewhat arbitrary. A correct implementation will give you a `diff` of zero."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Block Implementations\n",
    "<a id=part1_2></a>\n",
    "\n",
    "We'll now implement some `Block`s that will enable us to later build an MLP model of arbitrary depth, complete with automatic differentiation.\n",
    "\n",
    "For each block, you'll first implement the `forward()` function.\n",
    "Then, you will calculate the derivative of the block by hand with respect to each of its\n",
    "input tensors and each of its parameter tensors (if any).\n",
    "Using your manually-calculated derivation, you can then implement the `backward()` function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Linear (fully connected) layer\n",
    "\n",
    "First, we'll implement an affine transform layer, also known as a fully connected layer.\n",
    "\n",
    "Given an input $\\vec{x}$ the layer computes,\n",
    "\n",
    "$$\n",
    "\\vec{z} = \\vec{x} \\mattr{W}  + \\vec{b} ,~\n",
    "\\mat{W}\\in\\set{R}^{D_{\\mathrm{out}}\\times D_{\\mathrm{in}}},~ \\vec{b}\\in\\set{R}^{D_{\\mathrm{out}}}.\n",
    "$$\n",
    "\n",
    "Notes:\n",
    "- We write it this way to follow the implementation conventions.\n",
    "- In the code, the input $\\vec{x}$ will always be a tensor containing a batch dimension first. Thanks to broadcasting, $\\vec{b}$ can remain a vector even if the input is a matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**TODO**: Complete the implementation of the `Linear` class in the `hw2/blocks.py` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 100\n",
    "in_features = 200\n",
    "num_classes = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Comparing gradients... \n",
      "input    diff=0.000\n",
      "param#01 diff=0.000\n",
      "param#02 diff=0.000\n"
     ]
    }
   ],
   "source": [
    "# Test Linear\n",
    "fc = blocks.Linear(in_features, 1000)\n",
    "x_test = torch.randn(N, in_features)\n",
    "\n",
    "# Test forward pass\n",
    "z = fc(x_test)\n",
    "test.assertSequenceEqual(z.shape, [N, 1000])\n",
    "\n",
    "# Test backward pass\n",
    "test_block_grad(fc, x_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Activation functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ReLU\n",
    "\n",
    "ReLU, or rectified linear unit is a very common activation function in deep learning architectures.\n",
    "In it's most standard form, as we'll implement here, it has no parameters.\n",
    "\n",
    "The ReLU operation is defined as\n",
    "\n",
    "$$\n",
    "\\mathrm{relu}(\\vec{x}) = \\max(\\vec{0},\\vec{x})\n",
    "$$\n",
    "\n",
    "Note that it's not strictly differentiable, however it has sub-gradients. The gradients w.r.t. $\\vec{x}$ are simply $1$ for any positive-valued elements and zero for negative-valued elements."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**TODO**: Complete the implementation of the `ReLU` class in the `hw2/blocks.py` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Comparing gradients... \n",
      "input    diff=0.000\n"
     ]
    }
   ],
   "source": [
    "# Test ReLU\n",
    "relu = blocks.ReLU()\n",
    "x_test = torch.randn(N, in_features)\n",
    "\n",
    "# Test forward pass\n",
    "z = relu(x_test)\n",
    "test.assertSequenceEqual(z.shape, x_test.shape)\n",
    "\n",
    "# Test backward pass\n",
    "test_block_grad(relu, x_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Sigmoid\n",
    "\n",
    "The sigmoid function $\\sigma(x)$ is also sometimes used as an activation function.\n",
    "We have also seen it previously in the context of logistic regression.\n",
    "\n",
    "The sigmoid function is defined as\n",
    "\n",
    "$$\n",
    "\\sigma(\\vec{x}) = \\frac{1}{1+\\exp(-\\vec{x})}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Comparing gradients... \n",
      "input    diff=0.000\n"
     ]
    }
   ],
   "source": [
    "# Test Sigmoid\n",
    "sigmoid = blocks.Sigmoid()\n",
    "x_test = torch.randn(N, in_features)\n",
    "\n",
    "# Test forward pass\n",
    "z = sigmoid(x_test)\n",
    "test.assertSequenceEqual(z.shape, x_test.shape)\n",
    "\n",
    "# Test backward pass\n",
    "test_block_grad(sigmoid, x_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cross-Entropy Loss\n",
    "\n",
    "As you know by know, this is a common loss function for classification tasks.\n",
    "In class, we defined it as \n",
    "\n",
    "$$\\ell_{\\mathrm{CE}}(\\vec{y},\\hat{\\vec{y}}) = - {\\vectr{y}} \\log(\\hat{\\vec{y}})$$\n",
    "\n",
    "where $\\hat{\\vec{y}} = \\mathrm{softmax}(x)$ is a probability vector (the output of softmax on the class scores $\\vec{x}$) and the vector $\\vec{y}$ is a 1-hot encoded class label.\n",
    "\n",
    "However, it's tricky to compute the gradient of softmax, so instead we'll define a version of cross-entropy that produces the exact same output but works directly on the class scores $\\vec{x}$.\n",
    "\n",
    "We can write,\n",
    "$$\\begin{align}\n",
    "\\ell_{\\mathrm{CE}}(\\vec{y},\\hat{\\vec{y}}) &= - {\\vectr{y}} \\log(\\hat{\\vec{y}}) \n",
    "= - {\\vectr{y}} \\log\\left(\\mathrm{softmax}(\\vec{x})\\right) \\\\\n",
    "&= - {\\vectr{y}} \\log\\left(\\frac{e^{\\vec{x}}}{\\sum_k e^{x_k}}\\right) \\\\\n",
    "&= - \\log\\left(\\frac{e^{x_y}}{\\sum_k e^{x_k}}\\right) \\\\\n",
    "&= - \\left(\\log\\left(e^{x_y}\\right) - \\log\\left(\\sum_k e^{x_k}\\right)\\right)\\\\\n",
    "&= - x_y + \\log\\left(\\sum_k e^{x_k}\\right)\n",
    "\\end{align}$$\n",
    "\n",
    "Where the scalar $y$ is the correct class label, so $x_y$ is the correct class score.\n",
    "\n",
    "Note that this version of cross entroy is also what's [provided](https://pytorch.org/docs/stable/nn.html#crossentropyloss) by PyTorch's `nn` module."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**TODO**: Complete the implementation of the `CrossEntropyLoss` class in the `hw2/blocks.py` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss= 2.7283620834350586\n",
      "Comparing gradients... \n",
      "input    diff=0.000\n"
     ]
    }
   ],
   "source": [
    "# Test CrossEntropy\n",
    "cross_entropy = blocks.CrossEntropyLoss()\n",
    "scores = torch.randn(N, num_classes)\n",
    "labels = torch.randint(low=0, high=num_classes, size=(N,), dtype=torch.long)\n",
    "\n",
    "# Test forward pass\n",
    "loss = cross_entropy(scores, labels)\n",
    "expected_loss = torch.nn.functional.cross_entropy(scores, labels)\n",
    "test.assertLess(torch.abs(expected_loss-loss).item(), 1e-5)\n",
    "print('loss=', loss.item())\n",
    "\n",
    "# Test backward pass\n",
    "test_block_grad(cross_entropy, scores, y=labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building Models\n",
    "<a id=part1_3></a>\n",
    "\n",
    "Now that we have some building `Block`s, we can build an MLP model of arbitrary depth and compute end-to-end gradients.\n",
    "\n",
    "First, lets copy an idea from PyTorch and implement our own version of the `nn.Sequential` `Module`.\n",
    "This is a `Block` which contains other `Block`s and calls them in sequence. We'll use this to build our MLP model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**TODO**: Complete the implementation of the `Sequential` class in the `hw2/blocks.py` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Comparing gradients... \n",
      "input    diff=0.000\n",
      "param#01 diff=0.000\n",
      "param#02 diff=0.000\n",
      "param#03 diff=0.000\n",
      "param#04 diff=0.000\n",
      "param#05 diff=0.000\n",
      "param#06 diff=0.000\n",
      "param#07 diff=0.000\n",
      "param#08 diff=0.000\n",
      "param#09 diff=0.000\n",
      "param#10 diff=0.000\n",
      "param#11 diff=0.000\n",
      "param#12 diff=0.000\n",
      "param#13 diff=0.000\n",
      "param#14 diff=0.000\n"
     ]
    }
   ],
   "source": [
    "# Test Sequential\n",
    "# Let's create a long sequence of blocks and see\n",
    "# whether we can compute end-to-end gradients of the whole thing.\n",
    "\n",
    "seq = blocks.Sequential(\n",
    "    blocks.Linear(in_features, 100),\n",
    "    blocks.Linear(100, 200),\n",
    "    blocks.Linear(200, 100),\n",
    "    blocks.ReLU(),\n",
    "    blocks.Linear(100, 500),\n",
    "    blocks.Linear(500, 200),\n",
    "    blocks.ReLU(),\n",
    "    blocks.Linear(200, 500),\n",
    "    blocks.ReLU(),\n",
    "    blocks.Linear(500, 1),\n",
    "    blocks.Sigmoid(),\n",
    ")\n",
    "x_test = torch.randn(N, in_features)\n",
    "\n",
    "# Test forward pass\n",
    "z = seq(x_test)\n",
    "test.assertSequenceEqual(z.shape, [N, 1])\n",
    "\n",
    "# Test backward pass\n",
    "test_block_grad(seq, x_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, equipped with a `Sequential`, all we have to do is create an MLP architecture.\n",
    "We'll define our MLP with the following hyperparameters:\n",
    "- Number of input features, $D$.\n",
    "- Number of output classes, $C$.\n",
    "- Sizes of hidden layers, $h_1,\\dots,h_L$.\n",
    "\n",
    "So the architecture will be:\n",
    "\n",
    "FC($D$, $h_1$) $\\rightarrow$ ReLU $\\rightarrow$\n",
    "FC($h_1$, $h_2$) $\\rightarrow$ ReLU $\\rightarrow$\n",
    "$\\cdots$ $\\rightarrow$\n",
    "FC($h_{L-1}$, $h_L$) $\\rightarrow$ ReLU $\\rightarrow$\n",
    "FC($h_{L}$, $C$)\n",
    "\n",
    "We'll also create a sequence of the above MLP and a cross-entropy loss, since it's the gradient of the loss that we need in order to train a model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**TODO**: Complete the implementation of the `MLP` class in the `hw2/models.py` module. Ignore the `dropout` parameter for now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MLP, Sequential\n",
      "\t[0] Linear(200, 100)\n",
      "\t[1] ReLU\n",
      "\t[2] Dropout(p=0)\n",
      "\t[3] Linear(100, 50)\n",
      "\t[4] ReLU\n",
      "\t[5] Dropout(p=0)\n",
      "\t[6] Linear(50, 100)\n",
      "\t[7] ReLU\n",
      "\t[8] Dropout(p=0)\n",
      "\t[9] Linear(100, 10)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import hw2.models as models\n",
    "\n",
    "# Create MLP model\n",
    "mlp = models.MLP(in_features, num_classes, hidden_features=[100, 50, 100])\n",
    "print(mlp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MLP loss=3.3796074390411377, activation=relu\n",
      "Comparing gradients... \n",
      "input    diff=0.000\n",
      "param#01 diff=0.000\n",
      "param#02 diff=0.000\n",
      "param#03 diff=0.000\n",
      "param#04 diff=0.000\n",
      "param#05 diff=0.000\n",
      "param#06 diff=0.000\n",
      "param#07 diff=0.000\n",
      "param#08 diff=0.000\n",
      "param#09 diff=0.000\n",
      "param#10 diff=0.000\n",
      "param#11 diff=0.000\n",
      "param#12 diff=0.000\n",
      "param#13 diff=0.000\n",
      "param#14 diff=0.000\n",
      "param#15 diff=0.000\n",
      "param#16 diff=0.000\n",
      "param#17 diff=0.000\n",
      "param#18 diff=0.000\n",
      "param#19 diff=0.000\n",
      "param#20 diff=0.000\n",
      "param#21 diff=0.000\n",
      "param#22 diff=0.000\n",
      "param#23 diff=0.000\n",
      "param#24 diff=0.000\n",
      "MLP loss=2.493561029434204, activation=sigmoid\n",
      "Comparing gradients... \n",
      "input    diff=0.000\n",
      "param#01 diff=0.000\n",
      "param#02 diff=0.000\n",
      "param#03 diff=0.000\n",
      "param#04 diff=0.000\n",
      "param#05 diff=0.000\n",
      "param#06 diff=0.000\n",
      "param#07 diff=0.000\n",
      "param#08 diff=0.000\n",
      "param#09 diff=0.000\n",
      "param#10 diff=0.000\n",
      "param#11 diff=0.000\n",
      "param#12 diff=0.000\n",
      "param#13 diff=0.000\n",
      "param#14 diff=0.000\n",
      "param#15 diff=0.000\n",
      "param#16 diff=0.000\n",
      "param#17 diff=0.000\n",
      "param#18 diff=0.000\n",
      "param#19 diff=0.000\n",
      "param#20 diff=0.000\n",
      "param#21 diff=0.000\n",
      "param#22 diff=0.000\n",
      "param#23 diff=0.000\n",
      "param#24 diff=0.000\n"
     ]
    }
   ],
   "source": [
    "# Test MLP architecture\n",
    "N = 100\n",
    "in_features = 10\n",
    "num_classes = 10\n",
    "for activation in ('relu', 'sigmoid'):\n",
    "    mlp = models.MLP(in_features, num_classes, hidden_features=[100, 50, 100], activation=activation)\n",
    "    test.assertEqual(len(mlp.sequence), 7)\n",
    "    \n",
    "    num_linear = 0\n",
    "    for b1, b2 in zip(mlp.sequence, mlp.sequence[1:]):\n",
    "        if (str(b2).lower() == activation):\n",
    "            test.assertTrue(str(b1).startswith('Linear'))\n",
    "            num_linear += 1\n",
    "            \n",
    "    test.assertTrue(str(mlp.sequence[-1]).startswith('Linear'))\n",
    "    test.assertEqual(num_linear, 3)\n",
    "\n",
    "    # Test MLP gradients\n",
    "    # Test forward pass\n",
    "    x_test = torch.randn(N, in_features)\n",
    "    labels = torch.randint(low=0, high=num_classes, size=(N,), dtype=torch.long)\n",
    "    z = mlp(x_test)\n",
    "    test.assertSequenceEqual(z.shape, [N, num_classes])\n",
    "\n",
    "    # Create a sequence of MLPs and CE loss\n",
    "    # Note: deliberately using the same MLP instance multiple times to create a recurrence.\n",
    "    seq_mlp = blocks.Sequential(mlp, mlp, mlp, blocks.CrossEntropyLoss())\n",
    "    loss = seq_mlp(x_test, y=labels)\n",
    "    test.assertEqual(loss.dim(), 0)\n",
    "    print(f'MLP loss={loss}, activation={activation}')\n",
    "\n",
    "    # Test backward pass\n",
    "    test_block_grad(seq_mlp, x_test, y=labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the above tests passed then congratulations - you've now implemented an arbitrarily deep model and loss function with end-to-end automatic differentiation!"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}