live_notes/lesson-093-minibatch-training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Returning to Implementation of a basic CNN\n",
    "\n",
    "<img src=\"https://snag.gy/yjkZ0C.jpg\" >\n",
    "\n",
    "# Minibatch Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai_lib import get_data, test_near\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import matplotlib as mpl\n",
    "\n",
    "mpl.rcParams['image.cmap'] = 'gray'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loading the dataset\n",
    "x_train, y_train, x_valid, y_valid = get_data()\n",
    "\n",
    "# getting some basic params\n",
    "n, m = x_train.shape\n",
    "n_classes = y_train.max() + 1\n",
    "n_hidden = 50\n",
    "\n",
    "# our simple model from before\n",
    "\n",
    "class Model(nn.Module):\n",
    "    def __init__(self, n_in, n_hidden, n_out):\n",
    "        super().__init__()\n",
    "        self.layers = [\n",
    "            nn.Linear(n_in, n_hidden),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(n_hidden, n_out)\n",
    "        ]\n",
    "    \n",
    "    def __call__(self, x):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model(m, n_hidden, 10)\n",
    "pred = model(x_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Review of Cross Entropy Loss\n",
    "\n",
    "This is based on softmax\n",
    "\n",
    "First, we will need to compute the softmax of our activations. This is defined by:\n",
    "\n",
    "$$\\hbox{softmax(x)}_{i} = \\frac{e^{x_{i}}}{e^{x_{0}} + e^{x_{1}} + \\cdots + e^{x_{n-1}}}$$\n",
    "\n",
    "or more concisely:\n",
    "\n",
    "$$\\hbox{softmax(x)}_{i} = \\frac{e^{x_{i}}}{\\sum_{0 \\leq j \\leq n-1} e^{x_{j}}}$$ \n",
    "\n",
    "In practice, we will need the log of the softmax when we calculate the loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>label</th>\n",
       "      <th>output</th>\n",
       "      <th>exp</th>\n",
       "      <th>softmapx</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>cat</td>\n",
       "      <td>0.882275</td>\n",
       "      <td>-0.125251</td>\n",
       "      <td>0.031979</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>dog</td>\n",
       "      <td>0.502540</td>\n",
       "      <td>-0.688080</td>\n",
       "      <td>0.175678</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>plane</td>\n",
       "      <td>0.244549</td>\n",
       "      <td>-1.408342</td>\n",
       "      <td>0.359574</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>fish</td>\n",
       "      <td>0.672065</td>\n",
       "      <td>-0.397400</td>\n",
       "      <td>0.101463</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>building</td>\n",
       "      <td>0.273180</td>\n",
       "      <td>-1.297626</td>\n",
       "      <td>0.331306</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      label    output       exp  softmapx\n",
       "0       cat  0.882275 -0.125251  0.031979\n",
       "1       dog  0.502540 -0.688080  0.175678\n",
       "2     plane  0.244549 -1.408342  0.359574\n",
       "3      fish  0.672065 -0.397400  0.101463\n",
       "4  building  0.273180 -1.297626  0.331306"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "df = pd.DataFrame({'label': ['cat', 'dog', 'plane', 'fish', 'building'], 'output': np.random.rand(5)})\n",
    "df['exp']= np.log(df['output'])\n",
    "df['softmapx'] = df['exp'] / df['exp'].sum()\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_softmax(x):\n",
    "    return (x.exp() / (x.exp().sum(-1, keepdim=True))).log()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cross entropy loss for some target $x$ and some prediction $p(x)$ is given by:\n",
    "\n",
    "$$ -\\sum x\\, \\log p(x) $$\n",
    "\n",
    "But since our $x$s are 1-hot encoded, this can be rewritten as $-\\log(p_{i})$ where i is the index of the desired target.\n",
    "\n",
    "<img src='https://snag.gy/jQnH5O.jpg' style=\"width:400px\" >"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_pred = log_softmax(pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([5, 0, 4])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([-2.3949, -2.2743, -2.3425], grad_fn=<IndexBackward>)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# consider the first 3 elements of the truth\n",
    "print(y_train[:3])\n",
    "\n",
    "# compare to our predictions\n",
    "# first 3\n",
    "row_idx = [0, 1, 2]\n",
    "\n",
    "# what was the probabilty for each class?\n",
    "col_idx = [5, 0 ,4]\n",
    "\n",
    "# whats the probability associated with 5, 0 and 4 \n",
    "# an example of advanced indexing\n",
    "sm_pred[row_idx, col_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def nll(inputs, target):\n",
    "    \"\"\" negative log-likelihood \"\"\" \n",
    "    return -inputs[range(target.shape[0]), target].mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.3053, grad_fn=<NegBackward>)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = nll(sm_pred, y_train)\n",
    "loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the formula \n",
    "\n",
    "$$\\log \\left ( \\frac{a}{b} \\right ) = \\log(a) - \\log(b)$$ \n",
    "\n",
    "gives a simplification when we compute the log softmax, which was previously defined as `(x.exp()/(x.exp().sum(-1,keepdim=True))).log()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_softmax(x):\n",
    "    return (x.exp() / (x.exp().sum(-1, keepdim=True))).log()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_near(nll(log_softmax(pred), y_train), loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A trick for stability sake \n",
    "\n",
    "Then, there is a way to compute the log of the sum of exponentials in a more stable way, called the [LogSumExp trick](https://en.wikipedia.org/wiki/LogSumExp). The idea is to use the following formula:\n",
    "\n",
    "$$\\log \\left ( \\sum_{j=1}^{n} e^{x_{j}} \\right ) = \\log \\left ( e^{a} \\sum_{j=1}^{n} e^{x_{j}-a} \\right ) = a + \\log \\left ( \\sum_{j=1}^{n} e^{x_{j}-a} \\right )$$\n",
    "\n",
    "where a is the maximum of the $x_{j}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def logsumexp(x):\n",
    "    # find max\n",
    "    m = x.max(-1)[0]\n",
    "    \n",
    "    # and subtract and do summation of logs\n",
    "    return m + (x-m[:,None]).exp().sum(-1).log()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_near(logsumexp(pred), pred.logsumexp(-1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "so now we can replace and update our `log_softmax` function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_softmax(x):\n",
    "    return x - x.logsumexp(-1, keepdim=True)\n",
    "\n",
    "test_near(nll(log_softmax(pred), y_train), loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The pytorch version\n",
    "test_near(F.nll_loss(F.log_softmax(pred, -1), y_train), loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the above function is combined into cross entropy\n",
    "test_near(F.cross_entropy(pred, y_train), loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_func = F.cross_entropy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def accuracy(out, y_true):\n",
    "    \"\"\"\n",
    "    torch.argmax(out, dim=1) == yb <-- returns bools\n",
    "    .float() <-- changes to float\n",
    "    .mean() <-- gives average accuracy\n",
    "    \"\"\"\n",
    "    return (torch.argmax(out, dim=1) == y_true).float().mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([ 0.0158,  0.0640, -0.0864, -0.1477,  0.0200, -0.0735,  0.1565,  0.1732,\n",
       "          0.0614, -0.0440], grad_fn=<SelectBackward>), torch.Size([64, 10]))"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 64 # batchsize\n",
    "lr = 0.5 # learning rate\n",
    "epochs = 1 # how many epoches to train for\n",
    "\n",
    "n_batches = (n-1) // bs + 1\n",
    "\n",
    "xb = x_train[0:bs]\n",
    "yb = y_train[0:bs]\n",
    "preds = model(xb)\n",
    "\n",
    "# look at the predictions (reminder tha the model has not been trained yet)\n",
    "preds[0], preds.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.3138, grad_fn=<NllLossBackward>)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss_func(preds, yb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0469)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "accuracy(preds, yb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Loop!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(epochs):\n",
    "    for i in range(n_batches):\n",
    "        \n",
    "        # get indices of batch\n",
    "        start_i = i*bs\n",
    "        end_i = start_i + bs\n",
    "        \n",
    "        # pull our batch out\n",
    "        xb = x_train[start_i: end_i]\n",
    "        yb = y_train[start_i: end_i]\n",
    "        \n",
    "        # calculate our loss\n",
    "        loss = loss_func(model(xb), yb)\n",
    "        \n",
    "        # backpropogation\n",
    "        loss.backward()\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            \n",
    "            # go through all layers\n",
    "            for layer in model.layers:\n",
    "                \n",
    "                # if weights are found, update the weights\n",
    "                # and update the bias\n",
    "                if hasattr(layer, 'weight'):\n",
    "                    layer.weight -= layer.weight.grad * lr\n",
    "                    layer.bias -= layer.bias.grad * lr\n",
    "                    \n",
    "                    # zero out the parameters for the next\n",
    "                    # iteration\n",
    "                    layer.weight.grad.zero_()\n",
    "                    layer.bias.grad.zero_()\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.1419, grad_fn=<NllLossBackward>), tensor(0.9375))"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss_func(model(xb), yb), accuracy(model(xb), yb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simplify the training loop\n",
    "\n",
    "### Parameters\n",
    "\n",
    "Use nn.Module.__setattr__ and move relu to functional. \n",
    "\n",
    "So instead of looping through LAYERS, we will be looping through the parameters. Note the `DummyModule` below. Using the `parameters` class method, if we can collect all the layer-related attributes in this class as modules, we can easily loop through all the parameters.\n",
    "\n",
    "#### How is modules self-populated?\n",
    "\n",
    "There's a special `__setattr__` will automatically be called anytime the `self.some_attr` is being called. In the section below, everytime a new layer is set, referenace is added to the dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DummyModule():\n",
    "    def __init__(self, n_in, n_hidden, n_out):\n",
    "        self._modules = {}\n",
    "        self.l1 = nn.Linear(n_in, n_hidden)\n",
    "        self.l2 = nn.Linear(n_hidden, n_out)\n",
    "    \n",
    "    def __setattr__(self, k, v):\n",
    "        if not k.startswith(\"_\"):\n",
    "            self._modules[k] = v\n",
    "        super().__setattr__(k, v)\n",
    "        \n",
    "    def __repr__(self):\n",
    "        return f'{self._modules}'\n",
    "        \n",
    "    def parameters(self):\n",
    "        for l in self._modules.values():\n",
    "            for p in l.parameters():\n",
    "                yield p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'l1': Linear(in_features=784, out_features=50, bias=True), 'l2': Linear(in_features=50, out_features=10, bias=True)}"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mdl = DummyModule(m, n_hidden, 10)\n",
    "mdl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Example of looping through the parameters\n",
    "\n",
    "Even though we didnt create the layers in list or dictionary, they are still called and displayed using the `self.parameters`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[torch.Size([50, 784]),\n",
       " torch.Size([50]),\n",
       " torch.Size([10, 50]),\n",
       " torch.Size([10])]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[o.shape for o in mdl.parameters()]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pytorch has a similar setattr method\n",
    "\n",
    "Will use `nn.Module.__setattr__` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self, n_in, n_hidden, n_out):\n",
    "        super().__init__()\n",
    "        self.l1 = nn.Linear(n_in, n_hidden)\n",
    "        self.l2 = nn.Linear(n_hidden, n_out)\n",
    "        \n",
    "    def __call__(self, x):\n",
    "        x = self.l1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.l2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Note that we can see all the layers as follows\n",
    "\n",
    "Not that we can easily pull out any of the name nn.layer objects using the `named_children` method which behaves a lot like the `parameters`method that we set above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "l1: Linear(in_features=784, out_features=50, bias=True)\n",
      "l2: Linear(in_features=50, out_features=10, bias=True)\n"
     ]
    }
   ],
   "source": [
    "model = Model(m, n_hidden, 10)\n",
    "for name, l in model.named_children():\n",
    "    print(f\"{name}: {l}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Also note that pytorch has a similar __repr__ already implemented"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Model(\n",
       "  (l1): Linear(in_features=784, out_features=50, bias=True)\n",
       "  (l2): Linear(in_features=50, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Linear(in_features=784, out_features=50, bias=True)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.l1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit():\n",
    "    for epoch in range(epochs):\n",
    "        for i in range(n_batches):\n",
    "\n",
    "            # get indices of batch\n",
    "            start_i = i*bs\n",
    "            end_i = start_i + bs\n",
    "\n",
    "            # pull our batch out\n",
    "            xb = x_train[start_i: end_i]\n",
    "            yb = y_train[start_i: end_i]\n",
    "\n",
    "            # calculate our loss\n",
    "            loss = loss_func(model(xb), yb)\n",
    "\n",
    "            # backpropogation\n",
    "            loss.backward()\n",
    "\n",
    "            with torch.no_grad():\n",
    "                for p in model.parameters():\n",
    "                    p -= p.grad * lr\n",
    "                    model.zero_grad()\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "cross_entropy() missing 1 required positional argument: 'target'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-34-462bf89083e8>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mfit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mloss_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m: cross_entropy() missing 1 required positional argument: 'target'"
     ]
    }
   ],
   "source": [
    "fit()\n",
    "loss_func(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Making the layers a list instead of individual attributes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# given a list of layers\n",
    "n_hidden = 50\n",
    "layers = [nn.Linear(m, n_hidden), nn.ReLU(), nn.Linear(n_hidden, 10)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self, layers):\n",
    "        super().__init__()\n",
    "        \n",
    "        # this by itself will not make them parameters\n",
    "        self.layers = layers\n",
    "        \n",
    "        # need to go through them and call \n",
    "        # add module, and will register all the layers\n",
    "        for idx, layer in enumerate(self.layers):\n",
    "            self.add_module(f'layer_{idx}', layer)\n",
    "                            \n",
    "    def __call__(self, x):\n",
    "        for layer in self.layers:\n",
    "            x =layer(x)\n",
    "        return x\n",
    "                            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model(layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Model(\n",
       "  (layer_0): Linear(in_features=784, out_features=50, bias=True)\n",
       "  (layer_1): ReLU()\n",
       "  (layer_2): Linear(in_features=50, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pytorch Module List\n",
    "\n",
    "This simplifies it even more, removing the loop registration, and using `ModuleList`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SequentialModel(nn.Module):\n",
    "    def __init__(self, layers):\n",
    "        super().__init__()\n",
    "        \n",
    "        # this by itself will not make them parameters\n",
    "        self.layers = nn.ModuleList(layers)\n",
    "                            \n",
    "    def __call__(self, x):\n",
    "        for layer in self.layers:\n",
    "            x =layer(x)\n",
    "        return x\n",
    "                            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SequentialModel(layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SequentialModel(\n",
       "  (layers): ModuleList(\n",
       "    (0): Linear(in_features=784, out_features=50, bias=True)\n",
       "    (1): ReLU()\n",
       "    (2): Linear(in_features=50, out_features=10, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## And even simplier still, nn.Sequential\n",
    "\n",
    "`nn.Sequential` will take care of that. If you look at the source code of Sequential, its almost exactly the same"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = nn.Sequential(nn.Linear(m, n_hidden), nn.ReLU(), nn.Linear(n_hidden, 10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mInit signature:\u001b[0m \u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mSequential\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\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;32mclass\u001b[0m \u001b[0mSequential\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mModule\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;34mr\"\"\"A sequential container.\u001b[0m\n",
       "\u001b[0;34m    Modules will be added to it in the order they are passed in the constructor.\u001b[0m\n",
       "\u001b[0;34m    Alternatively, an ordered dict of modules can also be passed in.\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m    To make it easier to understand, here is a small example::\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m        # Example of using Sequential\u001b[0m\n",
       "\u001b[0;34m        model = nn.Sequential(\u001b[0m\n",
       "\u001b[0;34m                  nn.Conv2d(1,20,5),\u001b[0m\n",
       "\u001b[0;34m                  nn.ReLU(),\u001b[0m\n",
       "\u001b[0;34m                  nn.Conv2d(20,64,5),\u001b[0m\n",
       "\u001b[0;34m                  nn.ReLU()\u001b[0m\n",
       "\u001b[0;34m                )\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m        # Example of using Sequential with OrderedDict\u001b[0m\n",
       "\u001b[0;34m        model = nn.Sequential(OrderedDict([\u001b[0m\n",
       "\u001b[0;34m                  ('conv1', nn.Conv2d(1,20,5)),\u001b[0m\n",
       "\u001b[0;34m                  ('relu1', nn.ReLU()),\u001b[0m\n",
       "\u001b[0;34m                  ('conv2', nn.Conv2d(20,64,5)),\u001b[0m\n",
       "\u001b[0;34m                  ('relu2', nn.ReLU())\u001b[0m\n",
       "\u001b[0;34m                ]))\u001b[0m\n",
       "\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;32mdef\u001b[0m \u001b[0m__init__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mSequential\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__init__\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;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m1\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mOrderedDict\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;32mfor\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitems\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[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd_module\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m            \u001b[0;32mfor\u001b[0m \u001b[0midx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd_module\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule\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;32mdef\u001b[0m \u001b[0m_get_item_by_idx\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0miterator\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midx\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\"\"\"Get the idx-th item of the iterator\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0msize\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0midx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moperator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mindex\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0msize\u001b[0m \u001b[0;34m<=\u001b[0m \u001b[0midx\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0msize\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m            \u001b[0;32mraise\u001b[0m \u001b[0mIndexError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'index {} is out of range'\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0midx\u001b[0m \u001b[0;34m%=\u001b[0m \u001b[0msize\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0;32mreturn\u001b[0m \u001b[0mnext\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mislice\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0miterator\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midx\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\n",
       "\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0;32mdef\u001b[0m \u001b[0m__getitem__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midx\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;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mslice\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;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__class__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mOrderedDict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_modules\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitems\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0midx\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;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m            \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_item_by_idx\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_modules\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midx\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;32mdef\u001b[0m \u001b[0m__setitem__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mkey\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_item_by_idx\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_modules\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mkeys\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0;32mreturn\u001b[0m \u001b[0msetattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule\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;32mdef\u001b[0m \u001b[0m__delitem__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midx\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;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mslice\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;32mfor\u001b[0m \u001b[0mkey\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_modules\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mkeys\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                \u001b[0mdelattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m            \u001b[0mkey\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_item_by_idx\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_modules\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mkeys\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m            \u001b[0mdelattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkey\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;32mdef\u001b[0m \u001b[0m__len__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\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;32mreturn\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_modules\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;32mdef\u001b[0m \u001b[0m__dir__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mkeys\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mSequential\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__dir__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mkeys\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mkey\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mkey\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mkeys\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0misdigit\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;32mreturn\u001b[0m \u001b[0mkeys\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;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\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;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_modules\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\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[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mFile:\u001b[0m           ~/envs/py3/lib/python3.7/site-packages/torch/nn/modules/container.py\n",
       "\u001b[0;31mType:\u001b[0m           type\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "nn.Sequential??"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### Let's return to our fit function\n",
    " \n",
    " Let's replace our previous manually coded optimization step:\n",
    "\n",
    "```python\n",
    "with torch.no_grad():\n",
    "    for p in model.parameters(): p -= p.grad * lr\n",
    "    model.zero_grad()\n",
    "```\n",
    "\n",
    "and instead use just:\n",
    "\n",
    "```python\n",
    "opt.step()\n",
    "opt.zero_grad()\n",
    "```\n",
    "\n",
    "Let's try and make the optimizer from scratch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Optimizer():\n",
    "    def __init__(self,params, lr=0.5):\n",
    "        self.lr = lr\n",
    "        self.params = list(params)\n",
    "        \n",
    "    def step(self):\n",
    "        \"\"\"\n",
    "        This method replaces  the \n",
    "         for p in model.parameters(): p -= p.grad * lr\n",
    "        \"\"\"\n",
    "        with torch.no_grad():\n",
    "            for p in self.params:\n",
    "                p -= p.grad * lr\n",
    "                \n",
    "    def zero_grad(self):\n",
    "        \"\"\"\n",
    "        This method implements\n",
    "        model.zero_grad()\n",
    "        \n",
    "        model might be too general, and too drastic\n",
    "        \"\"\"\n",
    "        for p in self.params:\n",
    "            p.grad.data.zero_()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = nn.Sequential(\n",
    "    nn.Linear(m, n_hidden),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(n_hidden, 10)\n",
    ")\n",
    "\n",
    "opt = Optimizer(model.parameters())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### let's sub in our optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(epochs):\n",
    "    for i in range((n-1)//bs + 1):\n",
    "        start_i = i*bs\n",
    "        end_i = start_i+bs\n",
    "        xb = x_train[start_i:end_i]\n",
    "        yb = y_train[start_i:end_i]\n",
    "        pred = model(xb)\n",
    "        loss = loss_func(pred, yb)\n",
    "\n",
    "        loss.backward()\n",
    "        opt.step()\n",
    "        opt.zero_grad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.2493, grad_fn=<NllLossBackward>), tensor(0.9375))"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss,acc = loss_func(model(xb), yb), accuracy(model(xb), yb)\n",
    "loss,acc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Luckily Pytorch takes care of this with `optim`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.3602, grad_fn=<NllLossBackward>)"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#export\n",
    "from torch import optim\n",
    "\n",
    "def get_model():\n",
    "    model = nn.Sequential(nn.Linear(m,n_hidden), nn.ReLU(), nn.Linear(n_hidden,10))\n",
    "    return model, optim.SGD(model.parameters(), lr=lr)\n",
    "\n",
    "model,opt = get_model()\n",
    "loss_func(model(xb), yb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(epochs):\n",
    "    for i in range((n-1)//bs + 1):\n",
    "        start_i = i*bs\n",
    "        end_i = start_i+bs\n",
    "        xb = x_train[start_i:end_i]\n",
    "        yb = y_train[start_i:end_i]\n",
    "        pred = model(xb)\n",
    "        loss = loss_func(pred, yb)\n",
    "\n",
    "        loss.backward()\n",
    "        opt.step()\n",
    "        opt.zero_grad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.1134, grad_fn=<NllLossBackward>), tensor(0.9375))"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss,acc = loss_func(model(xb), yb), accuracy(model(xb), yb)\n",
    "loss,acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert acc > 0.7"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Making Datasets and DataLoaders\n",
    "\n",
    "It's clunky to iterate through minibatches of x and y values separately:\n",
    "\n",
    "```python\n",
    "    xb = x_train[start_i:end_i]\n",
    "    yb = y_train[start_i:end_i]\n",
    "```\n",
    "\n",
    "Instead, let's do these two steps together, by introducing a `Dataset` class:\n",
    "\n",
    "```python\n",
    "    xb,yb = train_ds[i*bs : i*bs+bs]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dataset():\n",
    "    def __init__(self, x, y):\n",
    "        self.x = x\n",
    "        self.y = y\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.x)\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        return self.x[i], self.y[i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = Dataset(x_train, y_train)\n",
    "valid_ds = Dataset(x_valid, y_valid)\n",
    "\n",
    "assert len(train_ds)==len(x_train)\n",
    "assert len(valid_ds)==len(x_valid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
       "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "         [0., 0., 0.,  ..., 0., 0., 0.]]), tensor([5, 0, 4, 1, 9]))"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xb,yb = train_ds[0:5]\n",
    "assert xb.shape==(5,28*28)\n",
    "assert yb.shape==(5,)\n",
    "xb,yb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "model,opt = get_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(epochs):\n",
    "    for i in range((n-1)//bs + 1):\n",
    "        \n",
    "        # we have replace some of the other code\n",
    "        # with our dataset class\n",
    "        xb,yb = train_ds[i*bs : i*bs+bs]\n",
    "        pred = model(xb)\n",
    "        loss = loss_func(pred, yb)\n",
    "\n",
    "        loss.backward()\n",
    "        opt.step()\n",
    "        opt.zero_grad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.2123, grad_fn=<NllLossBackward>), tensor(0.9375))"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss,acc = loss_func(model(xb), yb), accuracy(model(xb), yb)\n",
    "assert acc>0.7\n",
    "loss,acc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Making a Dataloader\n",
    "\n",
    "The dataloader will manage things like:\n",
    "\n",
    "- batch loading\n",
    "- shuffling\n",
    "- indexes\n",
    "\n",
    "Previously, our loop iterated over batches (xb, yb) like this:\n",
    "\n",
    "```python\n",
    "for i in range((n-1)//bs + 1):\n",
    "    xb,yb = train_ds[i*bs : i*bs+bs]\n",
    "    ...\n",
    "```\n",
    "\n",
    "Let's make our loop much cleaner, using a data loader:\n",
    "\n",
    "```python\n",
    "for xb,yb in train_dl:\n",
    "    ...\n",
    "```\n",
    "\n",
    "`yield` is a great command to get familiar with, it will execute lazily and will only provide a value at a time, without loading the whole list into memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataLoader():\n",
    "    def __init__(self, ds, bs): \n",
    "        self.ds = ds\n",
    "        self.bs = bs\n",
    "        \n",
    "    def __iter__(self):\n",
    "        # to replace this with itertools\n",
    "        for i in range(0, len(self.ds), self.bs): \n",
    "            yield self.ds[i:i+self.bs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dl = DataLoader(train_ds, bs)\n",
    "valid_dl = DataLoader(valid_ds, bs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "xb,yb = next(iter(valid_dl))\n",
    "\n",
    "# check that these things are the right size\n",
    "assert xb.shape==(bs,28*28)\n",
    "assert yb.shape==(bs,)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's take a look at a sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(3)"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAP8AAAD8CAYAAAC4nHJkAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAADXZJREFUeJzt3X+MFPUZx/HPo2IEIYqQIuAp9jBNjPGgXkxNSaO2GGtIUBOJJjZXID1NNKlaTc3VpEbShDT1R+MfGIwErFatoJEoFiwhBbQx4o8qCiIaBM47rogRiBqLPv3j5uyJt99ddmd39njer+Ryu/PszDzZ8GFm9ju3X3N3AYjnmKIbAFAMwg8ERfiBoAg/EBThB4Ii/EBQhB8IivADQRF+IKjjGrkzM+N2QqDO3N0qeV1NR34zu9TM3jWz7WZ2ey3bAtBYVu29/WZ2rKRtkmZK2i3pFUnXuPs7iXU48gN11ogj//mStrv7B+7+paTHJc2uYXsAGqiW8E+WtGvQ893Zsm8xs04z22Rmm2rYF4Cc1f0DP3dfLGmxxGk/0ExqOfJ3S2oZ9Py0bBmAYaCW8L8i6SwzO9PMjpd0taSV+bQFoN6qPu1390NmdqOk1ZKOlbTE3d/OrTMAdVX1UF9VO+OaH6i7htzkA2D4IvxAUIQfCIrwA0ERfiAowg8ERfiBoAg/EBThB4Ii/EBQhB8IivADQRF+ICjCDwRF+IGgCD8QFOEHgiL8QFCEHwiK8ANBEX4gqIZO0Y3qtLW1Jes333xzyVpra2ty3VGjRiXrXV1dyfpJJ52UrD///PMlawcOHEiui/riyA8ERfiBoAg/EBThB4Ii/EBQhB8IivADQdU0S6+Z7ZB0QNJXkg65e3uZ1zNL7xBGjx6drO/cuTNZP/nkk/NsJ1fd3d0la6n7EyRp+fLlebcTQqWz9OZxk89F7r43h+0AaCBO+4Ggag2/S1pjZq+aWWceDQFojFpP+2e4e7eZfU/SC2a21d3XD35B9p8C/zEATaamI7+7d2e/+yQ9Len8IV6z2N3by30YCKCxqg6/mZ1oZmMGHku6RNLmvBoDUF+1nPZPkPS0mQ1s56/u/vdcugJQdzWN8x/xzhjnH9KYMWOS9VWrViXrH3/8ccna66+/nlx3+vTpyfoZZ5yRrLe0tCTrI0eOLFnbs2dPct0LLrggWS+3flSVjvMz1AcERfiBoAg/EBThB4Ii/EBQhB8IiqE+1GT8+PHJ+m233VZVTZLmzp2brC9btixZj4qhPgBJhB8IivADQRF+ICjCDwRF+IGgCD8QFFN0oyZ796a/uPnFF18sWSs3zl/uz40Z568NR34gKMIPBEX4gaAIPxAU4QeCIvxAUIQfCIpxftRk7NixyXpXV1fV2540aVLV66I8jvxAUIQfCIrwA0ERfiAowg8ERfiBoAg/EFTZ7+03syWSZknqc/dzsmWnSHpC0hRJOyTNcfdPyu6M7+0fdtra2pL1J598MlmfOnVqydq2bduS686cOTNZ37VrV7IeVZ7f279U0qWHLbtd0lp3P0vS2uw5gGGkbPjdfb2kfYctni1p4GtUlkm6POe+ANRZtdf8E9y9J3vcK2lCTv0AaJCa7+13d09dy5tZp6TOWvcDIF/VHvn3mNlEScp+95V6obsvdvd2d2+vcl8A6qDa8K+U1JE97pD0TD7tAGiUsuE3s8ck/UvSD8xst5nNl7RQ0kwze0/Sz7LnAIaRsuP8ue6Mcf6m09HRkazfddddyXpLS0uy/vnnn5eszZo1K7nuunXrknUMLc9xfgBHIcIPBEX4gaAIPxAU4QeCIvxAUHx191Fg9OjRJWu33nprct077rgjWT/mmPTxYd++w//m69tmzJhRsrZ169bkuqgvjvxAUIQfCIrwA0ERfiAowg8ERfiBoAg/EBTj/EeBpUuXlqxdeeWVNW17+fLlyfp9992XrDOW37w48gNBEX4gKMIPBEX4gaAIPxAU4QeCIvxAUIzzHwVaW1vrtu1FixYl6y+99FLd9o364sgPBEX4gaAIPxAU4QeCIvxAUIQfCIrwA0GVHec3syWSZknqc/dzsmV3SvqVpP9kL+ty91X1ahJpa9asKVlra2ur27al8vcBLFy4sGTto48+qqon5KOSI/9SSZcOsfxed5+W/RB8YJgpG353Xy8pPS0LgGGnlmv+G83sTTNbYmZjc+sIQENUG/5FklolTZPUI+nuUi80s04z22Rmm6rcF4A6qCr87r7H3b9y968lPSjp/MRrF7t7u7u3V9skgPxVFX4zmzjo6RWSNufTDoBGqWSo7zFJF0oab2a7Jf1e0oVmNk2SS9oh6bo69gigDszdG7czs8btLJCRI0eWrD3yyCPJdc8777xk/fTTT6+qpwG9vb0la3Pnzk2uu3r16pr2HZW7WyWv4w4/ICjCDwRF+IGgCD8QFOEHgiL8QFAM9R3lTjjhhGT9uOPSt3rs378/z3a+5YsvvkjWb7nllmT9gQceyLOdowZDfQCSCD8QFOEHgiL8QFCEHwiK8ANBEX4gKMb5kXTuuecm6/fee2+yftFFF1W97507dybrU6ZMqXrbRzPG+QEkEX4gKMIPBEX4gaAIPxAU4QeCIvxAUIzzN4FRo0Yl65999lmDOjlyY8emp2lcsmRJydrs2bNr2vfkyZOT9Z6enpq2P1wxzg8gifADQRF+ICjCDwRF+IGgCD8QFOEHgkp/abskM2uR9LCkCZJc0mJ3/7OZnSLpCUlTJO2QNMfdP6lfq8NXa2trsr5x48Zk/bnnnkvWN2/eXLJWbqx7/vz5yfqIESOS9XJj7VOnTk3WU95///1kPeo4fl4qOfIfkvQbdz9b0o8k3WBmZ0u6XdJadz9L0trsOYBhomz43b3H3V/LHh+QtEXSZEmzJS3LXrZM0uX1ahJA/o7omt/MpkiaLullSRPcfeC8q1f9lwUAhomy1/wDzGy0pBWSbnL3/Wb/v33Y3b3Ufftm1imps9ZGAeSroiO/mY1Qf/AfdfenssV7zGxiVp8oqW+odd19sbu3u3t7Hg0DyEfZ8Fv/If4hSVvc/Z5BpZWSOrLHHZKeyb89APVSyWn/jyX9QtJbZvZGtqxL0kJJfzOz+ZI+lDSnPi0Of1dddVWyfuqppybr8+bNy7OdIzL48m4otfxJ+MGDB5P166+/vupto7yy4Xf3jZJK/Qv4ab7tAGgU7vADgiL8QFCEHwiK8ANBEX4gKMIPBFXx7b2o3rhx44puoW5WrFiRrC9YsKBkra9vyJtCv9Hb21tVT6gMR34gKMIPBEX4gaAIPxAU4QeCIvxAUIQfCIopuhug3NdfX3zxxcn6tddem6xPmjSpZO3TTz9NrlvO/fffn6xv2LAhWT906FBN+8eRY4puAEmEHwiK8ANBEX4gKMIPBEX4gaAIPxAU4/zAUYZxfgBJhB8IivADQRF+ICjCDwRF+IGgCD8QVNnwm1mLma0zs3fM7G0z+3W2/E4z6zazN7Kfy+rfLoC8lL3Jx8wmSpro7q+Z2RhJr0q6XNIcSQfd/U8V74ybfIC6q/Qmn7Iz9rh7j6Se7PEBM9siaXJt7QEo2hFd85vZFEnTJb2cLbrRzN40syVmNrbEOp1mtsnMNtXUKYBcVXxvv5mNlvRPSX9w96fMbIKkvZJc0gL1XxrMK7MNTvuBOqv0tL+i8JvZCEnPSlrt7vcMUZ8i6Vl3P6fMdgg/UGe5/WGPmZmkhyRtGRz87IPAAVdI2nykTQIoTiWf9s+QtEHSW5K+zhZ3SbpG0jT1n/bvkHRd9uFgalsc+YE6y/W0Py+EH6g//p4fQBLhB4Ii/EBQhB8IivADQRF+ICjCDwRF+IGgCD8QFOEHgiL8QFCEHwiK8ANBEX4gqLJf4JmzvZI+HPR8fLasGTVrb83al0Rv1cqztzMqfWFD/57/Ozs32+Tu7YU1kNCsvTVrXxK9Vauo3jjtB4Ii/EBQRYd/ccH7T2nW3pq1L4neqlVIb4Ve8wMoTtFHfgAFKST8Znapmb1rZtvN7PYieijFzHaY2VvZzMOFTjGWTYPWZ2abBy07xcxeMLP3st9DTpNWUG9NMXNzYmbpQt+7ZpvxuuGn/WZ2rKRtkmZK2i3pFUnXuPs7DW2kBDPbIand3QsfEzazn0g6KOnhgdmQzOyPkva5+8LsP86x7v7bJuntTh3hzM116q3UzNK/VIHvXZ4zXuehiCP/+ZK2u/sH7v6lpMclzS6gj6bn7usl7Tts8WxJy7LHy9T/j6fhSvTWFNy9x91fyx4fkDQws3Sh712ir0IUEf7JknYNer5bzTXlt0taY2avmlln0c0MYcKgmZF6JU0ospkhlJ25uZEOm1m6ad67ama8zhsf+H3XDHf/oaSfS7ohO71tSt5/zdZMwzWLJLWqfxq3Hkl3F9lMNrP0Ckk3ufv+wbUi37sh+irkfSsi/N2SWgY9Py1b1hTcvTv73SfpafVfpjSTPQOTpGa/+wru5xvuvsfdv3L3ryU9qALfu2xm6RWSHnX3p7LFhb93Q/VV1PtWRPhfkXSWmZ1pZsdLulrSygL6+A4zOzH7IEZmdqKkS9R8sw+vlNSRPe6Q9EyBvXxLs8zcXGpmaRX83jXdjNfu3vAfSZep/xP/9yX9rogeSvT1fUn/zn7eLro3SY+p/zTwv+r/bGS+pHGS1kp6T9I/JJ3SRL39Rf2zOb+p/qBNLKi3Geo/pX9T0hvZz2VFv3eJvgp537jDDwiKD/yAoAg/EBThB4Ii/EBQhB8IivADQRF+ICjCDwT1P0/hXGStUmRfAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.imshow(xb[0].view(28,28))\n",
    "yb[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit():\n",
    "    for epoch in range(epochs):\n",
    "        for xb,yb in train_dl:\n",
    "            pred = model(xb)\n",
    "            loss = loss_func(pred, yb)\n",
    "            loss.backward()\n",
    "            opt.step()\n",
    "            opt.zero_grad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.0570, grad_fn=<NllLossBackward>), tensor(0.9844))"
      ]
     },
     "execution_count": 77,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss,acc = loss_func(model(xb), yb), accuracy(model(xb), yb)\n",
    "assert acc>0.7\n",
    "loss,acc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Random sampling\n",
    "\n",
    "We want our training set to be in a random order, and that order should differ each iteration. But the validation set shouldn't be randomized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sampler():\n",
    "    def __init__(self, ds, bs, shuffle=False):\n",
    "        self.n = len(ds)\n",
    "        self.bs = bs\n",
    "        self.shuffle = shuffle\n",
    "        \n",
    "    def __iter__(self):\n",
    "        if self.shuffle:\n",
    "            # if we are shuffling, we pull some random indexes\n",
    "            self.idxs = torch.randperm(self.n)\n",
    "        else:\n",
    "            # otherwise pull teh indexes in order\n",
    "            self.idxs = torch.arange(self.n)\n",
    "            \n",
    "        for i in range(0, self.n, self.bs): \n",
    "            yield self.idxs[i:i+self.bs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "small_ds = Dataset(*train_ds[:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6, 7, 8]), tensor([9])]"
      ]
     },
     "execution_count": 86,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s = Sampler(small_ds,3,False)\n",
    "[o for o in s]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([1, 5, 0]), tensor([2, 3, 9]), tensor([8, 4, 7]), tensor([6])]"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s = Sampler(small_ds,3,True)\n",
    "[o for o in s]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate(b):\n",
    "    \"\"\"\n",
    "    Will grabs the Xs and Ys and stacks them up\n",
    "    \"\"\"\n",
    "    xs,ys = zip(*b)\n",
    "    return torch.stack(xs),torch.stack(ys)\n",
    "\n",
    "class DataLoader():\n",
    "    def __init__(self, ds, sampler, collate_fn=collate):\n",
    "        self.ds = ds\n",
    "        self.sampler = sampler\n",
    "        self.collate_fn = collate_fn\n",
    "        \n",
    "    def __iter__(self):\n",
    "        for s in self.sampler:\n",
    "            yield self.collate_fn([self.ds[i] for i in s])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we create two samplers, one randomized, the other not"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_samp = Sampler(train_ds, bs, shuffle=True)\n",
    "valid_samp = Sampler(valid_ds, bs, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dl = DataLoader(train_ds, sampler=train_samp, collate_fn=collate)\n",
    "valid_dl = DataLoader(valid_ds, sampler=valid_samp, collate_fn=collate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's confirm that the dataloaders are shuffling the way we expect them to"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(3)"
      ]
     },
     "execution_count": 93,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAP8AAAD8CAYAAAC4nHJkAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAADXZJREFUeJzt3X+MFPUZx/HPo2IEIYqQIuAp9jBNjPGgXkxNSaO2GGtIUBOJJjZXID1NNKlaTc3VpEbShDT1R+MfGIwErFatoJEoFiwhBbQx4o8qCiIaBM47rogRiBqLPv3j5uyJt99ddmd39njer+Ryu/PszDzZ8GFm9ju3X3N3AYjnmKIbAFAMwg8ERfiBoAg/EBThB4Ii/EBQhB8IivADQRF+IKjjGrkzM+N2QqDO3N0qeV1NR34zu9TM3jWz7WZ2ey3bAtBYVu29/WZ2rKRtkmZK2i3pFUnXuPs7iXU48gN11ogj//mStrv7B+7+paTHJc2uYXsAGqiW8E+WtGvQ893Zsm8xs04z22Rmm2rYF4Cc1f0DP3dfLGmxxGk/0ExqOfJ3S2oZ9Py0bBmAYaCW8L8i6SwzO9PMjpd0taSV+bQFoN6qPu1390NmdqOk1ZKOlbTE3d/OrTMAdVX1UF9VO+OaH6i7htzkA2D4IvxAUIQfCIrwA0ERfiAowg8ERfiBoAg/EBThB4Ii/EBQhB8IivADQRF+ICjCDwRF+IGgCD8QFOEHgiL8QFCEHwiK8ANBEX4gqIZO0Y3qtLW1Jes333xzyVpra2ty3VGjRiXrXV1dyfpJJ52UrD///PMlawcOHEiui/riyA8ERfiBoAg/EBThB4Ii/EBQhB8IivADQdU0S6+Z7ZB0QNJXkg65e3uZ1zNL7xBGjx6drO/cuTNZP/nkk/NsJ1fd3d0la6n7EyRp+fLlebcTQqWz9OZxk89F7r43h+0AaCBO+4Ggag2/S1pjZq+aWWceDQFojFpP+2e4e7eZfU/SC2a21d3XD35B9p8C/zEATaamI7+7d2e/+yQ9Len8IV6z2N3by30YCKCxqg6/mZ1oZmMGHku6RNLmvBoDUF+1nPZPkPS0mQ1s56/u/vdcugJQdzWN8x/xzhjnH9KYMWOS9VWrViXrH3/8ccna66+/nlx3+vTpyfoZZ5yRrLe0tCTrI0eOLFnbs2dPct0LLrggWS+3flSVjvMz1AcERfiBoAg/EBThB4Ii/EBQhB8IiqE+1GT8+PHJ+m233VZVTZLmzp2brC9btixZj4qhPgBJhB8IivADQRF+ICjCDwRF+IGgCD8QFFN0oyZ796a/uPnFF18sWSs3zl/uz40Z568NR34gKMIPBEX4gaAIPxAU4QeCIvxAUIQfCIpxftRk7NixyXpXV1fV2540aVLV66I8jvxAUIQfCIrwA0ERfiAowg8ERfiBoAg/EFTZ7+03syWSZknqc/dzsmWnSHpC0hRJOyTNcfdPyu6M7+0fdtra2pL1J598MlmfOnVqydq2bduS686cOTNZ37VrV7IeVZ7f279U0qWHLbtd0lp3P0vS2uw5gGGkbPjdfb2kfYctni1p4GtUlkm6POe+ANRZtdf8E9y9J3vcK2lCTv0AaJCa7+13d09dy5tZp6TOWvcDIF/VHvn3mNlEScp+95V6obsvdvd2d2+vcl8A6qDa8K+U1JE97pD0TD7tAGiUsuE3s8ck/UvSD8xst5nNl7RQ0kwze0/Sz7LnAIaRsuP8ue6Mcf6m09HRkazfddddyXpLS0uy/vnnn5eszZo1K7nuunXrknUMLc9xfgBHIcIPBEX4gaAIPxAU4QeCIvxAUHx191Fg9OjRJWu33nprct077rgjWT/mmPTxYd++w//m69tmzJhRsrZ169bkuqgvjvxAUIQfCIrwA0ERfiAowg8ERfiBoAg/EBTj/EeBpUuXlqxdeeWVNW17+fLlyfp9992XrDOW37w48gNBEX4gKMIPBEX4gaAIPxAU4QeCIvxAUIzzHwVaW1vrtu1FixYl6y+99FLd9o364sgPBEX4gaAIPxAU4QeCIvxAUIQfCIrwA0GVHec3syWSZknqc/dzsmV3SvqVpP9kL+ty91X1ahJpa9asKVlra2ur27al8vcBLFy4sGTto48+qqon5KOSI/9SSZcOsfxed5+W/RB8YJgpG353Xy8pPS0LgGGnlmv+G83sTTNbYmZjc+sIQENUG/5FklolTZPUI+nuUi80s04z22Rmm6rcF4A6qCr87r7H3b9y968lPSjp/MRrF7t7u7u3V9skgPxVFX4zmzjo6RWSNufTDoBGqWSo7zFJF0oab2a7Jf1e0oVmNk2SS9oh6bo69gigDszdG7czs8btLJCRI0eWrD3yyCPJdc8777xk/fTTT6+qpwG9vb0la3Pnzk2uu3r16pr2HZW7WyWv4w4/ICjCDwRF+IGgCD8QFOEHgiL8QFAM9R3lTjjhhGT9uOPSt3rs378/z3a+5YsvvkjWb7nllmT9gQceyLOdowZDfQCSCD8QFOEHgiL8QFCEHwiK8ANBEX4gKMb5kXTuuecm6/fee2+yftFFF1W97507dybrU6ZMqXrbRzPG+QEkEX4gKMIPBEX4gaAIPxAU4QeCIvxAUIzzN4FRo0Yl65999lmDOjlyY8emp2lcsmRJydrs2bNr2vfkyZOT9Z6enpq2P1wxzg8gifADQRF+ICjCDwRF+IGgCD8QFOEHgkp/abskM2uR9LCkCZJc0mJ3/7OZnSLpCUlTJO2QNMfdP6lfq8NXa2trsr5x48Zk/bnnnkvWN2/eXLJWbqx7/vz5yfqIESOS9XJj7VOnTk3WU95///1kPeo4fl4qOfIfkvQbdz9b0o8k3WBmZ0u6XdJadz9L0trsOYBhomz43b3H3V/LHh+QtEXSZEmzJS3LXrZM0uX1ahJA/o7omt/MpkiaLullSRPcfeC8q1f9lwUAhomy1/wDzGy0pBWSbnL3/Wb/v33Y3b3Ufftm1imps9ZGAeSroiO/mY1Qf/AfdfenssV7zGxiVp8oqW+odd19sbu3u3t7Hg0DyEfZ8Fv/If4hSVvc/Z5BpZWSOrLHHZKeyb89APVSyWn/jyX9QtJbZvZGtqxL0kJJfzOz+ZI+lDSnPi0Of1dddVWyfuqppybr8+bNy7OdIzL48m4otfxJ+MGDB5P166+/vupto7yy4Xf3jZJK/Qv4ab7tAGgU7vADgiL8QFCEHwiK8ANBEX4gKMIPBFXx7b2o3rhx44puoW5WrFiRrC9YsKBkra9vyJtCv9Hb21tVT6gMR34gKMIPBEX4gaAIPxAU4QeCIvxAUIQfCIopuhug3NdfX3zxxcn6tddem6xPmjSpZO3TTz9NrlvO/fffn6xv2LAhWT906FBN+8eRY4puAEmEHwiK8ANBEX4gKMIPBEX4gaAIPxAU4/zAUYZxfgBJhB8IivADQRF+ICjCDwRF+IGgCD8QVNnwm1mLma0zs3fM7G0z+3W2/E4z6zazN7Kfy+rfLoC8lL3Jx8wmSpro7q+Z2RhJr0q6XNIcSQfd/U8V74ybfIC6q/Qmn7Iz9rh7j6Se7PEBM9siaXJt7QEo2hFd85vZFEnTJb2cLbrRzN40syVmNrbEOp1mtsnMNtXUKYBcVXxvv5mNlvRPSX9w96fMbIKkvZJc0gL1XxrMK7MNTvuBOqv0tL+i8JvZCEnPSlrt7vcMUZ8i6Vl3P6fMdgg/UGe5/WGPmZmkhyRtGRz87IPAAVdI2nykTQIoTiWf9s+QtEHSW5K+zhZ3SbpG0jT1n/bvkHRd9uFgalsc+YE6y/W0Py+EH6g//p4fQBLhB4Ii/EBQhB8IivADQRF+ICjCDwRF+IGgCD8QFOEHgiL8QFCEHwiK8ANBEX4gqLJf4JmzvZI+HPR8fLasGTVrb83al0Rv1cqztzMqfWFD/57/Ozs32+Tu7YU1kNCsvTVrXxK9Vauo3jjtB4Ii/EBQRYd/ccH7T2nW3pq1L4neqlVIb4Ve8wMoTtFHfgAFKST8Znapmb1rZtvN7PYieijFzHaY2VvZzMOFTjGWTYPWZ2abBy07xcxeMLP3st9DTpNWUG9NMXNzYmbpQt+7ZpvxuuGn/WZ2rKRtkmZK2i3pFUnXuPs7DW2kBDPbIand3QsfEzazn0g6KOnhgdmQzOyPkva5+8LsP86x7v7bJuntTh3hzM116q3UzNK/VIHvXZ4zXuehiCP/+ZK2u/sH7v6lpMclzS6gj6bn7usl7Tts8WxJy7LHy9T/j6fhSvTWFNy9x91fyx4fkDQws3Sh712ir0IUEf7JknYNer5bzTXlt0taY2avmlln0c0MYcKgmZF6JU0ospkhlJ25uZEOm1m6ad67ama8zhsf+H3XDHf/oaSfS7ohO71tSt5/zdZMwzWLJLWqfxq3Hkl3F9lMNrP0Ckk3ufv+wbUi37sh+irkfSsi/N2SWgY9Py1b1hTcvTv73SfpafVfpjSTPQOTpGa/+wru5xvuvsfdv3L3ryU9qALfu2xm6RWSHnX3p7LFhb93Q/VV1PtWRPhfkXSWmZ1pZsdLulrSygL6+A4zOzH7IEZmdqKkS9R8sw+vlNSRPe6Q9EyBvXxLs8zcXGpmaRX83jXdjNfu3vAfSZep/xP/9yX9rogeSvT1fUn/zn7eLro3SY+p/zTwv+r/bGS+pHGS1kp6T9I/JJ3SRL39Rf2zOb+p/qBNLKi3Geo/pX9T0hvZz2VFv3eJvgp537jDDwiKD/yAoAg/EBThB4Ii/EBQhB8IivADQRF+ICjCDwT1P0/hXGStUmRfAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "xb,yb = next(iter(valid_dl))\n",
    "plt.imshow(xb[0].view(28,28))\n",
    "yb[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(8)"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAP8AAAD8CAYAAAC4nHJkAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAADd9JREFUeJzt3X2oVHUex/HPd80lUiMf2MslY20rFsQi7WJCtrm0ZZlgQZhSYWx4e4S1LDL3jw2WhVwy6Y8QrmRqz1smSS2rrdjWQoQPuJa2mYWldq+357TH7frdP+5x95bOb8aZM3Pm+n2/YLgz5zvnnC+DH8+ZOQ8/c3cBiOcnRTcAoBiEHwiK8ANBEX4gKMIPBEX4gaAIPxAU4QeCIvxAUMc1cmVmxumEQJ25u1Xyvpq2/GZ2iZm9ZWY7zWxeLcsC0FhW7bn9ZjZA0g5JF0naI2mDpJnuvj0xD1t+oM4aseUfL2mnu7/r7t9JelLStBqWB6CBagn/yZJ293m9J5v2A2bWbmYbzWxjDesCkLO6/+Dn7h2SOiR2+4FmUsuWf6+kU/q8HplNA9AP1BL+DZLOMLNTzeynkmZIWp1PWwDqrerdfnf/3sxulbRG0gBJS919W26dAairqg/1VbUyvvMDddeQk3wA9F+EHwiK8ANBEX4gKMIPBEX4gaAIPxAU4QeCIvxAUIQfCIrwA0ERfiAowg8ERfiBoAg/EBThB4Ii/EBQhB8IivADQRF+ICjCDwRF+IGgCD8QFOEHgiL8QFCEHwiK8ANBEX4gKMIPBFX1EN2SZGa7JO2X1CPpe3dvy6MpAPVXU/gzv3b3j3JYDoAGYrcfCKrW8LuktWa2ycza82gIQGPUuts/0d33mtnPJL1oZv9295f7viH7T4H/GIAmY+6ez4LM7pF0wN3vS7wnn5UBKMndrZL3Vb3bb2aDzGzIoeeSLpb0RrXLA9BYtez2t0haZWaHlvO4u/8tl64A1F1uu/0VrYzd/mNOW1v61I4pU6ZUvez3338/WX/kkUeS9Z6enqrXXasxY8Yk648//njVyz7rrLOS9brv9gPo3wg/EBThB4Ii/EBQhB8IivADQeVxVR+OYfPmzUvWb7/99mR9xIgRebbzA5MmTUrWr7vuuqqXfeKJJybrc+fOTdZnz56drO/cubNk7cYbb0zOmxe2/EBQhB8IivADQRF+ICjCDwRF+IGgCD8QFMf5g/v888+T9cGDByfra9asSdaXLFlS9bwrV65M1keOHJmsDxgwoGTt3HPPTc67atWqZH348OHJ+m233ZasP/jggyVrBw8eTM6bF7b8QFCEHwiK8ANBEX4gKMIPBEX4gaAIPxAUx/mPcanj7JI0aNCgZD11PFqS5s+fn6wfOHAgWU+5++67k/UPPvggWU8dq586dWpy3k2bNiXr5513XrKeul6/WbDlB4Ii/EBQhB8IivADQRF+ICjCDwRF+IGgyh7nN7OlkqZK6nb3Mdm0YZKekjRK0i5J09390/q1iZQTTjihZG38+PHJeT/++ONkfcGCBcl6Lcfxyzn++OOT9WeeeSZZnzBhQsnazTffnJz34YcfTta//fbbZL0/qGTLv0zSJT+aNk/SOnc/Q9K67DWAfqRs+N39ZUmf/GjyNEnLs+fLJV2ec18A6qza7/wt7t6ZPe+S1JJTPwAapOZz+93dzcxL1c2sXVJ7resBkK9qt/z7zKxVkrK/3aXe6O4d7t7m7m1VrgtAHVQb/tWSZmXPZ0l6Lp92ADRK2fCb2ROSXpX0SzPbY2bXS7pX0kVm9rak32SvAfQjZb/zu/vMEqULc+4FVRo3blzJ2ujRo5Pz7tixI1n/7LPPqurpkNT9AlasWJGc97LLLkvWP/zww2R98uTJJWvr169PzhsBZ/gBQRF+ICjCDwRF+IGgCD8QFOEHguLW3ce41DDVknTSSScl60OGDKlp+a+88krJ2plnnpmcd9u2bcn6jBkzapo/Orb8QFCEHwiK8ANBEX4gKMIPBEX4gaAIPxCUuZe8A1f+K0vc7gvVGzp0aMna9u3bk/Pu378/WX/00UeT9ZtuuilZb2kpfXvHOXPmJOd9+umnk/XOzs5kPSp3t0rex5YfCIrwA0ERfiAowg8ERfiBoAg/EBThB4LiOP8xIDVE9wsvvJCc94ILLqhp3eVu7Z26dXi5W2/39PRU1VN0HOcHkET4gaAIPxAU4QeCIvxAUIQfCIrwA0GVvW+/mS2VNFVSt7uPyabdI2m2pEMHaue7+1/r1STSvvnmm5K17u7umpZdbv4JEyYk611dXTWtH/VTyZZ/maRLjjB9kbufnT0IPtDPlA2/u78s6ZMG9AKggWr5zn+rmW01s6VmVvo+UgCaUrXhXyzpNElnS+qUtLDUG82s3cw2mtnGKtcFoA6qCr+773P3Hnc/KGmJpPGJ93a4e5u7t1XbJID8VRV+M2vt8/IKSW/k0w6ARqnkUN8TkiZJGmFmeyT9QdIkMztbkkvaJemGOvYIoA64nr8fGDZsWLI+derUkrVly5bVtO7Nmzcn621tfJtrNlzPDyCJ8ANBEX4gKMIPBEX4gaAIPxBU2eP8qL9zzjknWV+8eHGynjrc9tJLLyXn/fLLL5P1Sy+9NFmfPHlysr5mzZpkHcVhyw8ERfiBoAg/EBThB4Ii/EBQhB8IivADQXGcvwHKXfa6aNGiZH3s2LHJ+l133VWy1tHRkZz39NNPT9bLHee/8sork/W1a9eWrDXycnIcji0/EBThB4Ii/EBQhB8IivADQRF+ICjCDwTFrbtzUO7W2uWO41977bU1zT937txkvRarVq1K1qdNm5asp85xKHdbcFSHW3cDSCL8QFCEHwiK8ANBEX4gKMIPBEX4gaDKXs9vZqdIWiGpRZJL6nD3B8xsmKSnJI2StEvSdHf/tH6tNq8HHnggWb/66quT9aVLlybrd9xxx1H3lJdy1/uXkxqTgOP8xapky/+9pLnuPlrSBEm3mNloSfMkrXP3MySty14D6CfKht/dO919c/Z8v6Q3JZ0saZqk5dnblku6vF5NAsjfUX3nN7NRksZKek1Si7t3ZqUu9X4tANBPVHwPPzMbLGmlpDnu/oXZ/08fdncvdd6+mbVLaq+1UQD5qmjLb2YD1Rv8x9z92WzyPjNrzeqtkrqPNK+7d7h7m7un72IJoKHKht96N/EPSXrT3e/vU1otaVb2fJak5/JvD0C9VLLbf56kayW9bmZbsmnzJd0r6S9mdr2k9yRNr0+Lze+aa65J1pctW5as33nnncl6PS+7HjhwYN2WLUm7d++u6/JRvbLhd/d/Sip1ffCF+bYDoFE4ww8IivADQRF+ICjCDwRF+IGgCD8QFLfuzsHevXuT9dbW1mR93bp1yXpXV9dR91Spq666Klk/7rj00eAtW7Yk6xMnTixZ++qrr5LzojrcuhtAEuEHgiL8QFCEHwiK8ANBEX4gKMIPBMVx/hyUG6J74cKFyfr555+fZztH5dVXX03WFyxYkKy/8847yfrXX3991D2hNhznB5BE+IGgCD8QFOEHgiL8QFCEHwiK8ANBcZwfOMZwnB9AEuEHgiL8QFCEHwiK8ANBEX4gKMIPBFU2/GZ2ipmtN7PtZrbNzH6XTb/HzPaa2ZbsMaX+7QLIS9mTfMysVVKru282syGSNkm6XNJ0SQfc/b6KV8ZJPkDdVXqST3o4lt4FdUrqzJ7vN7M3JZ1cW3sAinZU3/nNbJSksZJeyybdamZbzWypmQ0tMU+7mW00s401dQogVxWf229mgyX9Q9Kf3P1ZM2uR9JEkl/RH9X41+G2ZZbDbD9RZpbv9FYXfzAZKel7SGne//wj1UZKed/cxZZZD+IE6y+3CHjMzSQ9JerNv8LMfAg+5QtIbR9skgOJU8mv/REmvSHpd0sFs8nxJMyWdrd7d/l2Sbsh+HEwtiy0/UGe57vbnhfAD9cf1/ACSCD8QFOEHgiL8QFCEHwiK8ANBEX4gKMIPBEX4gaAIPxAU4QeCIvxAUIQfCIrwA0GVvYFnzj6S9F6f1yOyac2oWXtr1r4keqtWnr39vNI3NvR6/sNWbrbR3dsKayChWXtr1r4keqtWUb2x2w8ERfiBoIoOf0fB609p1t6atS+J3qpVSG+FfucHUJyit/wAClJI+M3sEjN7y8x2mtm8Inooxcx2mdnr2cjDhQ4xlg2D1m1mb/SZNszMXjSzt7O/RxwmraDemmLk5sTI0oV+ds024nXDd/vNbICkHZIukrRH0gZJM919e0MbKcHMdklqc/fCjwmb2a8kHZC04tBoSGb2Z0mfuPu92X+cQ939ribp7R4d5cjNdeqt1MjS16nAzy7PEa/zUMSWf7ykne7+rrt/J+lJSdMK6KPpufvLkj750eRpkpZnz5er9x9Pw5XorSm4e6e7b86e75d0aGTpQj+7RF+FKCL8J0va3ef1HjXXkN8uaa2ZbTKz9qKbOYKWPiMjdUlqKbKZIyg7cnMj/Whk6ab57KoZ8Tpv/OB3uInuPk7SpZJuyXZvm5L3fmdrpsM1iyWdpt5h3DolLSyymWxk6ZWS5rj7F31rRX52R+irkM+tiPDvlXRKn9cjs2lNwd33Zn+7Ja1S79eUZrLv0CCp2d/ugvv5H3ff5+497n5Q0hIV+NllI0uvlPSYuz+bTS78sztSX0V9bkWEf4OkM8zsVDP7qaQZklYX0MdhzGxQ9kOMzGyQpIvVfKMPr5Y0K3s+S9JzBfbyA80ycnOpkaVV8GfXdCNeu3vDH5KmqPcX/3ck/b6IHkr09QtJ/8oe24ruTdIT6t0N/I96fxu5XtJwSeskvS3p75KGNVFvj6h3NOet6g1aa0G9TVTvLv1WSVuyx5SiP7tEX4V8bpzhBwTFD35AUIQfCIrwA0ERfiAowg8ERfiBoAg/EBThB4L6L7/oeFswDRzHAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "xb,yb = next(iter(train_dl))\n",
    "plt.imshow(xb[0].view(28,28))\n",
    "yb[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0)"
      ]
     },
     "execution_count": 95,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAP8AAAD8CAYAAAC4nHJkAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAADkhJREFUeJzt3X+sFfWZx/HPgwX8Uf6QLYt40b0IuoaoiF7N/oGmumvDkibQxJCSSNgAi8GaLIl/VK1GYl3TrOKqMVRvAymQrmXjT0I2tCzZaE20AbSrovy4ktuUn1ehAZEoCk//uMPuBe/5zuGcOWcGn/crubnnzHNmzuPI587MmTPzNXcXgHiGlN0AgHIQfiAowg8ERfiBoAg/EBThB4Ii/EBQhB8IivADQX2rnW9mZnydEGgxd7d6XtfUlt/MpprZNjPrMbN7m1kWgPayRr/bb2bnSNou6TZJuyRtlDTL3T9IzMOWH2ixdmz5b5TU4+473f2YpF9Lmt7E8gC0UTPh75D0pwHPd2XTTmFmC8xsk5ltauK9ABSs5R/4uXu3pG6J3X6gSprZ8u+WdMmA52OzaQDOAs2Ef6Oky81snJkNk/RDSWuKaQtAqzW82+/uX5nZ3ZJ+I+kcScvdfUthnQFoqYZP9TX0ZhzzAy3Xli/5ADh7EX4gKMIPBEX4gaAIPxAU4QeCIvxAUIQfCIrwA0ERfiAowg8ERfiBoAg/EBThB4Ii/EBQhB8IivADQRF+ICjCDwRF+IGgCD8QFOEHgiL8QFCEHwiK8ANBEX4gKMIPBEX4gaAIPxBUw0N0S5KZ9Ur6VNJxSV+5e1cRTQH1GDlyZLI+fPjwhpd99OjRZP3QoUMNL7sqmgp/5hZ3/6SA5QBoI3b7gaCaDb9L+q2ZbTazBUU0BKA9mt3tn+Luu83sryWtN7Ot7v76wBdkfxT4wwBUTFNbfnffnf3uk/SypBsHeU23u3fxYSBQLQ2H38wuMLMRJx9L+p6k94tqDEBrNbPbP1rSy2Z2cjn/4e7rCukKQMuZu7fvzcza92ZoiyFD0juPl156ac1ab29vct6FCxcm6/fdd1+y3tHRkayn9PT0JOtPPvlksv7cc88l6ydOnDjjnurl7lbP6zjVBwRF+IGgCD8QFOEHgiL8QFCEHwiqiKv68A120UUXJeuzZ89O1mfNmlWzdt111yXnfeedd5L1t956K1mfPHlyzdpll12WnHfChAnJ+jPPPJOsHzlyJFlftWpVst4ObPmBoAg/EBThB4Ii/EBQhB8IivADQRF+ICjO8yNp3br0LRquueaaZH3jxo01a11d6Zs75Z3HnzlzZrI+atSomrXrr78+Oe+kSZOS9fnz5yfrV1xxRbJeBWz5gaAIPxAU4QeCIvxAUIQfCIrwA0ERfiAobt0dXGdnZ7K+ZcuWZP28885L1lP/vg4cOJCc9+abb07Wt27dmqy30sUXX5ysP/LII8n63Llzi2znFNy6G0AS4QeCIvxAUIQfCIrwA0ERfiAowg8ElXs9v5ktl/R9SX3uflU2baSk1ZI6JfVKmunuf25dm2jUjBkzkvWnn346WT/33HOLbOcUS5YsSdbLPI+fZ8+ePcl63vcAqqCeLf8vJU09bdq9kja4++WSNmTPAZxFcsPv7q9LOnja5OmSVmSPV0hKb14AVE6jx/yj3X1v9nifpNEF9QOgTZq+h5+7e+o7+2a2QNKCZt8HQLEa3fLvN7MxkpT97qv1Qnfvdvcud0/frRFAWzUa/jWS5mSP50h6tZh2ALRLbvjN7HlJb0r6WzPbZWbzJP1M0m1mtkPSP2TPAZxFuJ7/G27KlCnJ+tVXX52sr1+/Plnfvn17st7XV/OIMPfe9ocPH07WyzRx4sRkffPmzcl63n0QmsH1/ACSCD8QFOEHgiL8QFCEHwiK8ANBMUR3BXR0dCTrecNJr1mzpmbtjTfeSM6bV1+5cmWy/uWXXybr3d3dNWtn86m8vFOgzz77bJHttARbfiAowg8ERfiBoAg/EBThB4Ii/EBQhB8Iikt622DhwoXJ+uOPP56sDx06NFnv6qp9kySz9NWd69atS9Z7enqS9YcffjhZzzsfXlV5Q5MfOXIkWb/11luT9c8+++yMe6oXl/QCSCL8QFCEHwiK8ANBEX4gKMIPBEX4gaA4z1+nIUNq/5288847k/MuXbo0Wc/7f7Bz585kffXq1TVr8+fPT86bdz576tTTB2g+1bFjx5L1Mp1//vk1a2+++WZy3rz/rltuuSVZz/seQCtxnh9AEuEHgiL8QFCEHwiK8ANBEX4gKMIPBJV7nt/Mlkv6vqQ+d78qm7ZY0j9L+jh72f3u/l+5b3YWn+efMGFCzdq2bdva2MnXHThwoGZt1apVyXnvueeeottpm0mTJiXra9eurVkbPnx4ct677rorWX/hhReS9TIVeZ7/l5IG+6bHv7v7tdlPbvABVEtu+N39dUkH29ALgDZq5pj/bjN718yWm9mFhXUEoC0aDf/PJY2XdK2kvZKW1HqhmS0ws01mtqnB9wLQAg2F3933u/txdz8h6ReSbky8ttvdu9y99l0mAbRdQ+E3szEDnv5A0vvFtAOgXXKH6Daz5yV9V9J3zGyXpIckfdfMrpXkknolpa9pBVA5XM+fuf3225P1hx56qGYtbyz3PJ9//nmy/tprryXrd9xxR83awYPVPVEzbNiwZH3RokXJ+rx585L1jz76qGYttc6kaq+3PFzPDyCJ8ANBEX4gKMIPBEX4gaAIPxBU7nn+KPJuUd3M6bze3t5k/dFHH03Wly1b1vB7l+3KK6+sWXviiSeS8950003JeuqW5VL+bcujY8sPBEX4gaAIPxAU4QeCIvxAUIQfCIrwA0FxSW/mxIkTyXpqPe3Zsyc57w033JCs79u3L1mvsgcffDBZnzt3bs3aF198kZz3scceS9bP5u8/tBKX9AJIIvxAUIQfCIrwA0ERfiAowg8ERfiBoLieP7Nz585kfdy4cQ0ve/z48cn62LFjk/VRo0Y1/N55w1hPmzat4WVL6WGwJWn27Nk1azt27EjOu3///oZ6Qn3Y8gNBEX4gKMIPBEX4gaAIPxAU4QeCIvxAULnX85vZJZJWShotySV1u/tTZjZS0mpJnZJ6Jc109z/nLKuy1/PPmjUrWV+8eHHN2oQJE5LzHj16NFk/fvx4sj5ixIhkvRk9PT3J+tKlS5P1p556qsh2UIAir+f/StI97j5R0t9J+pGZTZR0r6QN7n65pA3ZcwBnidzwu/ted387e/yppA8ldUiaLmlF9rIVkma0qkkAxTujY34z65Q0WdLvJY12971ZaZ/6DwsAnCXq/m6/mX1b0ouSFrn7YbP/P6xwd691PG9mCyQtaLZRAMWqa8tvZkPVH/xfuftL2eT9ZjYmq4+R1DfYvO7e7e5d7t5VRMMAipEbfuvfxC+T9KG7DxxWdY2kOdnjOZJeLb49AK1Sz6m+KZJ+J+k9SSfvb32/+o/7/1PSpZL+qP5TfQdzllXZU315Ojs7a9YeeOCB9jUyiK1bt9asvfLKK8l5Dx06lKx//PHHDfWE8tR7qi/3mN/d35BUa2F/fyZNAagOvuEHBEX4gaAIPxAU4QeCIvxAUIQfCIohuoFvGIboBpBE+IGgCD8QFOEHgiL8QFCEHwiK8ANBEX4gKMIPBEX4gaAIPxAU4QeCIvxAUIQfCIrwA0ERfiAowg8ERfiBoAg/EBThB4Ii/EBQhB8IivADQeWG38wuMbP/MbMPzGyLmf1LNn2xme02sz9kP9Na3y6AouQO2mFmYySNcfe3zWyEpM2SZkiaKemIuz9e95sxaAfQcvUO2vGtOha0V9Le7PGnZvahpI7m2gNQtjM65jezTkmTJf0+m3S3mb1rZsvN7MIa8ywws01mtqmpTgEUqu6x+szs25Jek/Sv7v6SmY2W9Ikkl/RT9R8azM1ZBrv9QIvVu9tfV/jNbKiktZJ+4+5PDFLvlLTW3a/KWQ7hB1qssIE6zcwkLZP04cDgZx8EnvQDSe+faZMAylPPp/1TJP1O0nuSTmST75c0S9K16t/t75V0Z/bhYGpZbPmBFit0t78ohB9ovcJ2+wF8MxF+ICjCDwRF+IGgCD8QFOEHgiL8QFCEHwiK8ANBEX4gKMIPBEX4gaAIPxAU4QeCyr2BZ8E+kfTHAc+/k02roqr2VtW+JHprVJG9/U29L2zr9fxfe3OzTe7eVVoDCVXtrap9SfTWqLJ6Y7cfCIrwA0GVHf7ukt8/paq9VbUvid4aVUpvpR7zAyhP2Vt+ACUpJfxmNtXMtplZj5ndW0YPtZhZr5m9l408XOoQY9kwaH1m9v6AaSPNbL2Z7ch+DzpMWkm9VWLk5sTI0qWuu6qNeN323X4zO0fSdkm3SdolaaOkWe7+QVsbqcHMeiV1uXvp54TN7GZJRyStPDkakpn9m6SD7v6z7A/nhe7+44r0tlhnOHJzi3qrNbL0P6nEdVfkiNdFKGPLf6OkHnff6e7HJP1a0vQS+qg8d39d0sHTJk+XtCJ7vEL9/3jarkZvleDue9397ezxp5JOjixd6rpL9FWKMsLfIelPA57vUrWG/HZJvzWzzWa2oOxmBjF6wMhI+ySNLrOZQeSO3NxOp40sXZl118iI10XjA7+vm+Lu10n6R0k/ynZvK8n7j9mqdLrm55LGq38Yt72SlpTZTDay9IuSFrn74YG1MtfdIH2Vst7KCP9uSZcMeD42m1YJ7r47+90n6WX1H6ZUyf6Tg6Rmv/tK7uf/uPt+dz/u7ick/UIlrrtsZOkXJf3K3V/KJpe+7gbrq6z1Vkb4N0q63MzGmdkwST+UtKaEPr7GzC7IPoiRmV0g6Xuq3ujDayTNyR7PkfRqib2coiojN9caWVolr7vKjXjt7m3/kTRN/Z/4fyTpJ2X0UKOvyyT9b/azpezeJD2v/t3AL9X/2cg8SX8laYOkHZL+W9LICvW2Sv2jOb+r/qCNKam3KerfpX9X0h+yn2llr7tEX6WsN77hBwTFB35AUIQfCIrwA0ERfiAowg8ERfiBoAg/EBThB4L6C/GJt+DIBStjAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "xb,yb = next(iter(train_dl))\n",
    "plt.imshow(xb[0].view(28,28))\n",
    "yb[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.2036, grad_fn=<NllLossBackward>), tensor(0.9531))"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model,opt = get_model()\n",
    "fit()\n",
    "\n",
    "loss,acc = loss_func(model(xb), yb), accuracy(model(xb), yb)\n",
    "assert acc>0.7\n",
    "loss,acc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PyTorch DataLoader\n",
    "\n",
    "Let's take a look at the pre-built stuff, and we will see that the code does a very similar thing to what we have made"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from torch.utils.data import DataLoader, SequentialSampler, RandomSampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dl = DataLoader(train_ds, bs, sampler=RandomSampler(train_ds), collate_fn=collate)\n",
    "valid_dl = DataLoader(valid_ds, bs, sampler=SequentialSampler(valid_ds), collate_fn=collate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [],
   "source": [
    "xb,yb = next(iter(train_dl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.0851, grad_fn=<NllLossBackward>), tensor(0.9844))"
      ]
     },
     "execution_count": 100,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model,opt = get_model()\n",
    "fit()\n",
    "loss_func(model(xb), yb), accuracy(model(xb), yb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dl = DataLoader(train_ds, bs, shuffle=True, drop_last=True)\n",
    "valid_dl = DataLoader(valid_ds, bs, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [],
   "source": [
    "xb,yb = next(iter(train_dl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.3761, grad_fn=<NllLossBackward>), tensor(0.9375))"
      ]
     },
     "execution_count": 104,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model,opt = get_model()\n",
    "fit()\n",
    "\n",
    "loss,acc = loss_func(model(xb), yb), accuracy(model(xb), yb)\n",
    "assert acc>0.7\n",
    "loss,acc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that PyTorch's `DataLoader`, if you pass `num_workers`, will use multiple threads to call your `Dataset`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You **always** should also have a [validation set](http://www.fast.ai/2017/11/13/validation-sets/), in order to identify if you are overfitting.\n",
    "\n",
    "We will calculate and print the validation loss at the end of each epoch.\n",
    "\n",
    "(Note that we always call `model.train()` before training, and `model.eval()` before inference, because these are used by layers such as `nn.BatchNorm2d` and `nn.Dropout` to ensure appropriate behaviour for these different phases.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit(epochs, model, loss_func, opt, train_dl, valid_dl):\n",
    "    for epoch in range(epochs):\n",
    "        # Handle batchnorm / dropout\n",
    "        model.train()\n",
    "#         print(model.training)\n",
    "        for xb,yb in train_dl:\n",
    "            loss = loss_func(model(xb), yb)\n",
    "            loss.backward()\n",
    "            opt.step()\n",
    "            opt.zero_grad()\n",
    "\n",
    "        model.eval()\n",
    "#         print(model.training)\n",
    "        with torch.no_grad():\n",
    "            tot_loss,tot_acc = 0.,0.\n",
    "            for xb,yb in valid_dl:\n",
    "                pred = model(xb)\n",
    "                tot_loss += loss_func(pred, yb)\n",
    "                tot_acc  += accuracy (pred,yb)\n",
    "        nv = len(valid_dl)\n",
    "        print(epoch, tot_loss/nv, tot_acc/nv)\n",
    "    return tot_loss/nv, tot_acc/nv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Question*: Are these validation results correct if batch size varies?\n",
    "\n",
    "`get_dls` returns dataloaders for the training and validation sets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def get_dls(train_ds, valid_ds, bs, **kwargs):\n",
    "    return (DataLoader(train_ds, batch_size=bs, shuffle=True, **kwargs),\n",
    "            DataLoader(valid_ds, batch_size=bs*2, **kwargs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 tensor(0.1736) tensor(0.9481)\n",
      "1 tensor(0.1452) tensor(0.9576)\n",
      "2 tensor(0.1400) tensor(0.9571)\n",
      "3 tensor(0.1026) tensor(0.9688)\n",
      "4 tensor(0.1564) tensor(0.9538)\n"
     ]
    }
   ],
   "source": [
    "train_dl,valid_dl = get_dls(train_ds, valid_ds, bs)\n",
    "model,opt = get_model()\n",
    "loss,acc = fit(5, model, loss_func, opt, train_dl, valid_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert acc>0.9"
   ]
  },
  {
   "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
}