02_fully_connected.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The forward and backward passes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from exp.nb_01 import *\n",
    "\n",
    "def get_data():\n",
    "    path = datasets.download_data(MNIST_URL, ext='.gz')\n",
    "    with gzip.open(path, 'rb') as f:\n",
    "        ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding='latin-1')\n",
    "    return map(tensor, (x_train,y_train,x_valid,y_valid))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def normalize(x,m,s): return (x-m)/s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train,y_train,x_valid,y_valid = get_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.1304), tensor(0.3073))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_mean, train_std = x_train.mean(), x_train.std()\n",
    "train_mean, train_std"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We want the mean and std to be as close as possible to 0 and 1 respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train = normalize(x_train, train_mean, train_std)\n",
    "# NB: Use training, not validation mean for validation set\n",
    "x_valid = normalize(x_valid, train_mean, train_std)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(-7.6999e-06), tensor(1.))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_mean, train_std = x_train.mean(), x_train.std()\n",
    "train_mean, train_std"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def test_near_zero(a,tol=1e-3): assert a.abs()<tol, f\"Near zero: {a}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_near_zero(train_mean)\n",
    "test_near_zero(train_std-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(50000, 784, tensor(9))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n, m = x_train.shape\n",
    "c = y_train.max()\n",
    "n,m,c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Foundations version"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Basic architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# num hidden\n",
    "nh = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "#simplified kaiming init/he init\n",
    "w1 = torch.randn(m,nh)/math.sqrt(m)\n",
    "b1 = torch.randn(nh)\n",
    "w2 = torch.randn(nh,1)/math.sqrt(nh)\n",
    "b2 = torch.randn(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0356)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w1.std()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_near_zero(w1.mean())\n",
    "test_near_zero(w1.std()-1/math.sqrt(m))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(-0.0059), tensor(0.9924))"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# This should be ~ (0,1) (mean,std)...\n",
    "x_valid.mean(),x_valid.std()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lin(x, w, b): return x@w + b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = lin(x_valid, w1, b1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(-0.0159), tensor(1.4171))"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#...so should this, because we used kaiming init, which is designed to do this\n",
    "t.mean(),t.std()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be studying this in a lot of depth as to why this kaiming initalization was 'THE THING THAT MATTERS'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def relu(x): return x.clamp_min(0.) # clamp_min is basically max(0,x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "t = relu(lin(x_valid, w1, b1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.5527), tensor(0.8353))"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#...actually it really should be this!\n",
    "t.mean(),t.std()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Kaiming he pointed out that the xavier init does not account for the impact of relu activation. This is a BIG problem. if your varaince halves each layer & you have a massive deep network then by the end your gradients will have totally dissapeared which is totally unacceptable. So they do something super genius which is they replace the 1 at the top with a 2.  \n",
    "So we do the same thing."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From pytorch docs: `a: the negative slope of the rectifier used after this layer (0 for ReLU by default)`\n",
    "\n",
    "$$\\text{std} = \\sqrt{\\frac{2}{(1 + a^2) \\times \\text{fan_in}}}$$\n",
    "\n",
    "This was introduced in the paper that described the Imagenet-winning approach from *He et al*: [Delving Deep into Rectifiers](https://arxiv.org/abs/1502.01852), which was also the first paper that claimed \"super-human performance\" on Imagenet (and, most importantly, it introduced resnets!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# kaiming init/he init for relu\n",
    "w1 = torch.randn(m,nh)*math.sqrt(2/m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.0004), tensor(0.0505))"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w1.mean(), w1.std()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.6053), tensor(0.9495))"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t = relu(lin(x_valid, w1, b1))\n",
    "t.mean(), t.std()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result is now much closer. Its not perfect and it varies some but it is much better than before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from torch.nn import init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "w1 = torch.zeros(m,nh)\n",
    "init.kaiming_normal_(w1, mode='fan_out')#check docs for mode\n",
    "t = relu(lin(x_valid, w1, b1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "#init.kaiming_normal_??\n",
    "# Choosing ``'fan_in' preserves the magnitude of the variance of the weights in the\n",
    "# forward pass. Choosing 'fan_out' preserves the magnitudes in the backwards pass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(-8.7934e-05), tensor(0.0505))"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w1.mean(), w1.std()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.6188), tensor(0.9679))"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t.mean(), t.std()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([784, 50])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w1.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([50, 784])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch.nn\n",
    "torch.nn.Linear(m,nh).weight.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.nn.Linear.forward??"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.nn.functional.linear??"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Turns out pytorch does a matrix product with a transpose.\n",
    "So thats why we have to give the opposite information(fan-out) when we were trying to do it with our linear layer which doesn't have a transpose. This was just meant as to how to look into the pytorch soure code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.nn.Conv2d??"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.nn.modules.conv._ConvNd??"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "# what if..?\n",
    "def relu(x): return x.clamp_min(0.) - 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.2189), tensor(1.0031))"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# kaiming init / he init for relu\n",
    "w1 = torch.randn(m,nh)*math.sqrt(2./m )\n",
    "t1 = relu(lin(x_valid, w1, b1))\n",
    "t1.mean(),t1.std()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Defining a new version of relu got our mean very close to zero. It also helped a bit to the std(closer to 1 than before)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model(xb):\n",
    "    l1 = lin(xb, w1, b1)\n",
    "    l2 = relu(l1)\n",
    "    l3 = lin(l2, w2, b2)\n",
    "    return l3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "14.8 ms ± 1.86 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit -n 10 _ = model(x_valid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10000, 1])"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.Size([x_valid.shape[0], 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert model(x_valid).shape==torch.Size([x_valid.shape[0],1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loss function: MSE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10000, 1])"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(x_valid).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need `squeeze()` to get rid of that trailing (,1), in order to use `mse`. (Of course, `mse` is not a suitable loss function for multi-class classification; we'll use a better loss function soon. We'll use `mse` for now to keep things simple.)  \n",
    "output.squeeze() gets rid of all ones in your matrix but suppose if you have a btach size of 1 then it will get rid of too causing an error. Hence, its better to specify which dimension you wanna squeeze"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def mse(output, targ): return (output.squeeze(-1) - targ).pow(2).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_train, y_valid = y_train.float(), y_valid.float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([50000, 1])"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "preds = model(x_train)\n",
    "preds.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(19.5981)"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mse(preds, y_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gradients and backward pass  \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "y^ = mse(lin2(relu(lin1(x))), y)  \n",
    "DOING THE CHAIN RULE. WE'RE DOING EACH BIT SEPERATELY  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 3, 1])"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = torch.zeros(3,3,1); a.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 3])"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a.squeeze().shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 3, 3, 1]), torch.Size([3, 1, 3, 1]))"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a.unsqueeze(0).shape, a.unsqueeze(1).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mse_grad(inp, targ):\n",
    "    # grad of loss wrt the output of previous layer.\n",
    "    # we store the gradient(derivative) in .g of the prev layer(which is inp)\n",
    "    # so that input of mse(inp(parameter of this func)) is same as output of the previous layer(which is a Linear layer)\n",
    "    inp.g = 2* (inp.squeeze()- targ).unsqueeze(-1)/inp.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "def relu_grad(inp, out):\n",
    "    # grad of relu with respect to input activations\n",
    "    # relu is max(0, x) so its deviative wil be:\n",
    "    # 0 if x<0\n",
    "    # 1 if x>0\n",
    "    # since we're doing chain rule we gotta multiply it with the gradient of the next layer(which is inp.g in mse_grad)\n",
    "    inp.g = (inp>0).float() * out.g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lin_grad(inp, out, w, b):\n",
    "    # grad of matmul wrt the input is simply the matrix product with the transpose.\n",
    "    # we don't just need the gradient of the o/ps wrt the i/ps\n",
    "    # but also wrt the weights and the biases. That's why we have three .g's here    \n",
    "    inp.g = out.g @ w.t()###\n",
    "    w.g = (inp.unsqueeze(-1)* out.g.unsqueeze(1)).sum(0)\n",
    "    b.g = out.g.sum(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward_and_backward(inp, targ):\n",
    "    # forward pass\n",
    "    l1 = inp @ w1 + b1\n",
    "    l2 = relu(l1)\n",
    "    out = l2 @ w2 + b2\n",
    "    # we don't actually need the loss in the backward pass\n",
    "    loss = mse(out, targ)\n",
    "    \n",
    "    # backward pass\n",
    "    mse_grad(out, targ)\n",
    "    lin_grad(l2, out, w2, b2)\n",
    "    relu_grad(l1, l2)\n",
    "    lin_grad(inp, l1, w1, b1)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall time: 5.89 s\n"
     ]
    }
   ],
   "source": [
    "%time forward_and_backward(x_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "# saving for testing against the pytorch library\n",
    "w1g = w1.g.clone()\n",
    "w2g = w2.g.clone()\n",
    "b1g = b1.g.clone()\n",
    "b2g = b2.g.clone()\n",
    "ig = x_train.g.clone()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We cheat a little bit and use PyTorch autograd to check our results.  \n",
    "Pytorch stores the same thing in .grad instead of .g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "xt2 = x_train.clone().requires_grad_(True)\n",
    "w12 = w1.clone().requires_grad_(True)\n",
    "w22 = w2.clone().requires_grad_(True)\n",
    "b12 = b1.clone().requires_grad_(True)\n",
    "b22 = b2.clone().requires_grad_(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward(inp, targ):\n",
    "    # forward pass\n",
    "    l1 = inp @ w12 + b12\n",
    "    l2 = relu(l1)\n",
    "    out = l2 @ w22 + b22\n",
    "    # we don't actually need the loss in backward\n",
    "    return mse(out, targ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = forward(xt2, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss.backward()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Testing if our calculated gradients are same to the ones tat were caluclated by Pytorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "test_near(w22.grad, w2g)\n",
    "test_near(b22.grad, b2g)\n",
    "test_near(w12.grad, w1g)\n",
    "test_near(b12.grad, b1g)\n",
    "test_near(xt2.grad, ig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Refactor model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Layers as classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Relu():\n",
    "    # if a class has a __call__ then we can treat it as a function.\n",
    "    def __call__(self, inp):\n",
    "        self.inp = inp\n",
    "        self.out = inp.clamp_min(0.)-0.5\n",
    "        return self.out\n",
    "    \n",
    "    def backward(self): self.inp.g = (self.inp>0).float() *self.out.g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Lin():\n",
    "    def __init__(self, w, b): \n",
    "        self.w, self.b = w, b\n",
    "        \n",
    "    def __call__(self, inp):\n",
    "        self.inp = inp\n",
    "        self.out = inp @ self.w + self.b\n",
    "        return self.out\n",
    "    \n",
    "    def backward(self): \n",
    "        self.inp.g = self.out.g @ self.w.t()\n",
    "        # creating a giant outer product just to sum it up is inefficient.\n",
    "        self.w.g = (self.inp.unsqueeze(-1) * self.out.g.unsqueeze(1)).sum(0)\n",
    "        self.b.g = self.out.g.sum(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mse():\n",
    "    def __call__(self, inp, targ):\n",
    "        self.inp = inp\n",
    "        self.targ = targ\n",
    "        self.out = (inp.squeeze() - targ).pow(2).mean()\n",
    "    \n",
    "    def backward(self):\n",
    "        self.inp.g = 2 * (self.inp.squeeze()- self.targ).unsqueeze(-1)/ self.targ.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model():\n",
    "    def __init__(self, w1, b1, w2,  b2):\n",
    "        self.layers = [Lin(w1, b1), Relu(), Lin(w2, b2)]\n",
    "        self.loss = Mse()\n",
    "        \n",
    "    def __call__(self, x, targ):\n",
    "        for l in self.layers: x = l(x)\n",
    "        return self.loss(x, targ)\n",
    "    \n",
    "    def backward(self):\n",
    "        self.loss.backward()\n",
    "        for l in reversed(self.layers): l.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "w1.g, b1.g, w2.g, b2.g = [None]*4 # setting our gradients to None\n",
    "model = Model(w1, b1, w2, b2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall time: 72.7 ms\n"
     ]
    }
   ],
   "source": [
    "%time loss = model(x_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall time: 11.5 s\n"
     ]
    }
   ],
   "source": [
    "%time model.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "test_near(w2g, w2.g)\n",
    "test_near(b2g, b2.g)\n",
    "test_near(w1g, w1.g)\n",
    "test_near(b1g, b1.g)\n",
    "test_near(ig, x_train.g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([784, 50]),\n",
       " torch.Size([784, 50]),\n",
       " torch.Size([50, 1]),\n",
       " torch.Size([50, 1]))"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w1g.shape, w1.g.shape, w2g.shape, w2.g.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Refactoring it again!  \n",
    "### Module.forward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Module():\n",
    "    def __call__(self, *args):\n",
    "        self.args = args\n",
    "        self.out = self.forward(*args)\n",
    "        return self.out\n",
    "    \n",
    "    def forward(self): raise Exception('not implemented')\n",
    "    def backward(self): self.bwd(self.out, *self.args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Relu(Module):\n",
    "    def forward(self, inp): return inp.clamp_min(0.)- 0.5\n",
    "    def bwd(self, out, inp): inp.g = (inp>0).float() * out.g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Lin(Module):\n",
    "    def __init__(self, w, b): self.w, self.b = w, b\n",
    "    \n",
    "    def forward(self, inp): return inp @ self.w + self.b\n",
    "    \n",
    "    def bwd(self, out, inp): \n",
    "        inp.g = out.g @ self.w.t()\n",
    "        self.w.g = torch.einsum('ai,aj->ij', inp, out.g)\n",
    "        self.b.g = out.g.sum(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mse(Module):\n",
    "    def forward (self, inp, targ): return (inp.squeeze() - targ).pow(2).mean()\n",
    "    def bwd(self, out, inp, targ): inp.g = 2*(inp.squeeze()-targ).unsqueeze(-1) / targ.shape[0]    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model():\n",
    "    def __init__(self):\n",
    "        self.layers = [Lin(w1, b1), Relu(), Lin(w2, b2)]\n",
    "        self.loss = Mse()\n",
    "    \n",
    "    def __call__(self, x, targ):\n",
    "        for l in self.layers: x = l(x)\n",
    "        return self.loss(x, targ)\n",
    "    \n",
    "    def backward(self):\n",
    "        self.loss.backward()\n",
    "        for l in reversed(self.layers): l.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "w1.g,b1.g,w2.g,b2.g = [None]*4\n",
    "model = Model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall time: 87.1 ms\n"
     ]
    }
   ],
   "source": [
    "%time loss = model(x_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall time: 234 ms\n"
     ]
    }
   ],
   "source": [
    "%time model.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_near(w2g, w2.g)\n",
    "test_near(b2g, b2.g)\n",
    "test_near(w1g, w1.g)\n",
    "test_near(b1g, b1.g)\n",
    "test_near(ig, x_train.g)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Without einsum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Lin(Module):\n",
    "    def __init__(self, w, b): self.w,self.b = w,b\n",
    "        \n",
    "    def forward(self, inp): return inp@self.w + self.b\n",
    "    \n",
    "    def bwd(self, out, inp):\n",
    "        inp.g = out.g @ self.w.t()\n",
    "        self.w.g = inp.t() @ out.g\n",
    "        self.b.g = out.g.sum(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "metadata": {},
   "outputs": [],
   "source": [
    "w1.g,b1.g,w2.g,b2.g = [None]*4\n",
    "model = Model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall time: 77.5 ms\n"
     ]
    }
   ],
   "source": [
    "%time loss = model(x_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall time: 215 ms\n"
     ]
    }
   ],
   "source": [
    "%time model.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_near(w2g, w2.g)\n",
    "test_near(b2g, b2.g)\n",
    "test_near(w1g, w1.g)\n",
    "test_near(b1g, b1.g)\n",
    "test_near(ig, x_train.g)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### nn.Linear and nn.Module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from torch import nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self, n_in, nh, n_out):\n",
    "        super().__init__()\n",
    "        self.layers = [nn.Linear(n_in, nh), nn.ReLU(), nn.Linear(nh, n_out)]\n",
    "        self.loss = mse\n",
    "        \n",
    "    def __call__(self, x, targ):\n",
    "        for l in self.layers: x = l(x)\n",
    "        return self.loss(x, targ)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(784, 50)"
      ]
     },
     "execution_count": 156,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m, nh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model(m, nh, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall time: 131 ms\n"
     ]
    }
   ],
   "source": [
    "%time loss = model(x_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall time: 101 ms\n"
     ]
    }
   ],
   "source": [
    "%time loss.backward()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our forward pass is as fast as the one of Pytorch's but their backward pass is twice as fast than ours. Jeremy guess that the reason for that must be that they are only calculating the gradients that they need whereas we're calculating all of them."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converted 02_fully_connected.ipynb to exp\\nb_02.py\n"
     ]
    }
   ],
   "source": [
    "!python notebook2script.py 02_fully_connected.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}