lec_9_v3.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6c3923a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision.transforms.functional as TF\n",
    "import torch.nn.functional as F\n",
    "from datasets import load_dataset\n",
    "import matplotlib.pyplot as plt\n",
    "from operator import itemgetter\n",
    "from torch.utils.data import default_collate, DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3197e03b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset fashion_mnist (/Users/arun/.cache/huggingface/datasets/fashion_mnist/fashion_mnist/1.0.0/0a671f063342996f19779d38c0ab4abef9c64f757b35af8134b331c294d7ba48)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cdab483921274a5c907e272edd8220ed",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "name = 'fashion_mnist'\n",
    "x,y = 'image','label'\n",
    "dsd = load_dataset(name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7d75abd8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['image', 'label'],\n",
       "        num_rows: 60000\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['image', 'label'],\n",
       "        num_rows: 10000\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dsd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b6018d7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def inplace(f):\n",
    "    def _f(b):\n",
    "        f(b)\n",
    "        return b\n",
    "    return _f\n",
    "\n",
    "@inplace\n",
    "def transformi_(b):\n",
    "    b[x] = [torch.flatten(TF.to_tensor(ele)) for ele in b[x]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "75126ca1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dls(train_ds, valid_ds, batch_size, **kwargs):\n",
    "    return (\n",
    "        DataLoader(train_ds, batch_size=batch_size, shuffle=True, **kwargs),\n",
    "        DataLoader(valid_ds, batch_size=batch_size*2, shuffle=False, **kwargs),\n",
    "    )\n",
    "\n",
    "def collate_dict(ds):\n",
    "    get = itemgetter(*ds.features)\n",
    "    def _f(b):\n",
    "        return get(default_collate(b))\n",
    "    return _f\n",
    "\n",
    "class DataLoaders:\n",
    "    def __init__(self, *ds):\n",
    "        self.train,self.valid = ds[:2]\n",
    "    \n",
    "    @classmethod\n",
    "    def from_dd(cls, dd, batch_size, as_tuple=True, **kwargs):\n",
    "        f = collate_dict(dd['train'])\n",
    "        return cls(*get_dls(*dd.values(), batch_size=batch_size, collate_fn=f, **kwargs))  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "676af169",
   "metadata": {},
   "outputs": [],
   "source": [
    "tds = dsd.with_transform(transformi_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b9825c73",
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 1024\n",
    "dls = DataLoaders.from_dd(tds, batch_size=bs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d04dac69",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = iter(dls.train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c06c177e",
   "metadata": {},
   "outputs": [],
   "source": [
    "xb, yb = next(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "868d6467",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Sneaker'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD4CAYAAAAq5pAIAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAAsTAAALEwEAmpwYAAAQqklEQVR4nO3dfWyd5XnH8d/lt9hx4iTGJDFJeE0GZQVSZsK0oJWKtYOoW0AaE9GEmIaUaoIJNKaNtdKKNmliW0u3PyZEKKxpV2grFQraGG0UWiG2KMWgAIEMAjQLeVlMEhLHIX6/9ocPkxt8X8ecd7i/Hyk69rn8nHPlOefnc3zu535uc3cB+ORrqncDAGqDsAOZIOxAJgg7kAnCDmSipZZ31mZzvF2dtbxLICvDOqlRH7GZamWF3cyulfRPkpolfdPd741+vl2dutKuKecuAQS2+9ZkreS38WbWLOmfJV0n6WJJG8zs4lJvD0B1lfM3+xpJb7r72+4+Kul7ktZXpi0AlVZO2JdJemfa9/sK1/0SM9toZv1m1j+mkTLuDkA5ygn7TB8CfOjYW3ff5O597t7Xqjll3B2AcpQT9n2SVkz7frmkA+W1A6Baygn785JWmdl5ZtYm6SZJT1amLQCVVvLQm7uPm9ntkn6sqaG3h9391Yp1BqCiyhpnd/enJD1VoV4AVBGHywKZIOxAJgg7kAnCDmSCsAOZIOxAJgg7kAnCDmSCsAOZIOxAJgg7kAnCDmSCsAOZIOxAJgg7kAnCDmSCsAOZIOxAJgg7kAnCDmSCsAOZIOxAJgg7kAnCDmSCsAOZIOxAJgg7kAnCDmSCsAOZIOxAJspastnM9kg6IWlC0ri791WiKQCVV1bYCz7n7ocrcDsAqoi38UAmyg27S/qJmb1gZhtn+gEz22hm/WbWP6aRMu8OQKnKfRu/1t0PmNliSVvM7L/d/dnpP+DumyRtkqQu6/Yy7w9Aicp6ZXf3A4XLAUmPS1pTiaYAVF7JYTezTjOb/8HXkr4gaWelGgNQWeW8jV8i6XEz++B2HnH3pyvSFYCKKzns7v62pMsq2AuAKmLoDcgEYQcyQdiBTBB2IBOEHcgEYQcyQdiBTBB2IBOEHcgEYQcyQdiBTBB2IBOEHchEJU44iU+wNx64Iqx37WoN60v/8b8q2c4nhs2Zky5OTITb+vh4SffJKzuQCcIOZIKwA5kg7EAmCDuQCcIOZIKwA5lgnP3jYOp03WkeLLTT1Bxu+t7N8boeNj4Z1n9219fC+jUjdyVrrSfDTXWqJ/5/995XZAw/2m/RPqsBH6n9Umi8sgOZIOxAJgg7kAnCDmSCsAOZIOxAJgg7kAnG2T8OyhgTvu31XWH9b3efF9YX/PvisL72wJ+F9a7fGUjWxn50ZrjtX//xt8P6A8/8blif3PFaWK+nk793ZbJ2+NL4Nficv9pW0n0WfWU3s4fNbMDMdk67rtvMtpjZ7sLlopLuHUDNzOZt/LckXXvadXdL2uruqyRtLXwPoIEVDbu7Pyvp6GlXr5e0ufD1ZknXV7YtAJVW6gd0S9z9oCQVLpN/2JnZRjPrN7P+MdX+eGAAU6r+aby7b3L3Pnfva1Vwkj0AVVVq2A+ZWa8kFS7TH7kCaAilhv1JSbcUvr5F0hOVaQdAtZgXGcM1s0clXS2pR9IhSV+V9CNJP5B0tqS9km5099M/xPuQLuv2K+2a8jrGh9gVlyRrB74Sn2N8aLAjrPtY/HrQvXgwrJ8YSt/+xHh8253zh8P6v67+l7B+00N/mqyt+Jvqns/+1Pr4PAHz7tyXrI1NxucgaLrmnWRtu2/VoB+dcSJ/0YNq3H1DokRqgY8RDpcFMkHYgUwQdiAThB3IBGEHMlF06K2Suqzbr2z6rfQPFOslOjWwxb+3rLW82bzVPPVv85J4Gukbf35+WO+84HiyNvje3HDbRWcMhfW5bWNhvbvj/bD+R2c9l6w9dvjycNsdh5aF9bGx+DG941efSdbmNsWP5yMH0lNQJWnP4e6wvuKMY2F90Zz0fpv0+BTaJ+7qTda273xAg0P7Z7wBXtmBTBB2IBOEHcgEYQcyQdiBTBB2IBOEHcjEx+tU0tE4vE/Em47E9XJYS7wbj9/YF9YPXRUfX9C8KJ7q2TMvvfbxmt694ba7j8encx4YnBfW9+/pCet/OXBDsnZJ74Fw26uW/SKs7xmKx7r/btt1YT1SbOruxUv/N6yf2R4fv3BkpDNZa2+OpyUfWZQ+45M3p8foeWUHMkHYgUwQdiAThB3IBGEHMkHYgUwQdiATtR9nr+H8+emaL1wZ1ocuisdsD18a7KrL4jHZvmXx0sHv/iKerz4+Ep9aeO+L6XnfA4eXh9taPKSrlnhqtdq648dzZDQ9n37HeNzb5y54I6z/wVnbw/q6lelTLm85lZ4TLkn/cfTSsF7sdM/7318Y1pssvd8+27073PaRMy5K1iaDB4xXdiAThB3IBGEHMkHYgUwQdiAThB3IBGEHMtFQ89lHfzue97332qDdIuPB3hSPBzeNxjew8vsnkrXmB98Nt/3Pf7ggrE+ebA3rNh73Nn5G+tzuJ9vj8WCbjG976UUDYf3m5S+F9Ss60nPSP92W3qeS1NOcnvMtSSMen9N+23B6Lv72E/FjEs03l6T25vi+ezviYy+6Wk6F9cjo/PRj5sHDXfSV3cweNrMBM9s57bp7zGy/me0o/Fv3EfsFUGOzeRv/LUnXznD9N9x9deHfU5VtC0ClFQ27uz8r6WgNegFQReV8QHe7mb1ceJu/KPVDZrbRzPrNrH9M1VsvDUCs1LDfL+kCSaslHZT09dQPuvsmd+9z975WpU+UB6C6Sgq7ux9y9wl3n5T0oKQ1lW0LQKWVFHYzmz4/8AZJO1M/C6AxFB1nN7NHJV0tqcfM9kn6qqSrzWy1JJe0R9KXZnVv8+dq4or0mtx7bij2u2cy3edEPF5sRcbRfWn8ecLrt3YE1XPCbbsXHgvrxywe021uTv+/JaltTnpS+vkrj4TbfnHxy2F9bcdbYb1J8fEL74wvSNaeOXVWuO2bw0vD+uN7LwvrR4+l9+uZ3fEY//kL4v22Yu57Yf28OfGxFycm25O1sWiwXNKZL6bPSf/WyfRzpWjY3X3DDFc/VGw7AI2Fw2WBTBB2IBOEHcgEYQcyQdiBTNR0iutEm2nw7OAourZ42mBPsIzuwo54yuCRk+lTGkvS8Gg8zbRlXnrZ5LlF+m5piofOzjsnfcpjSVpcZPnf+S3p3npa4yGmYsM8P3v/V8L6yGS8345PpIcsXx2MT+e8fO6xsP4nK58J67/Wnt6vy4s88xc0RUOt0uGJ9DLZkvTY0KqwPjaZbuDyzjfDbX/8fFe66Okc8MoOZIKwA5kg7EAmCDuQCcIOZIKwA5kg7EAmajrO3nzkpBZt3pasj3b9Rrj9+/N6krXBzniq5ejSeCz8/HPjUyYv6UiPVw+Nx2fg+XzPrrB+Vms8XfLIRPqUyJK0rjO9tPGhibZw22Lj7DuGzw7r85rTY/yStLwtPVX0wPDCcNu3h9KPtyQ9/danwvr4O+kprs3vx1OeW4rUi5wFW+1H42MrWobTz9enJz4bbtuhn8d3nsArO5AJwg5kgrADmSDsQCYIO5AJwg5kgrADmTD3eHy6krqs26+0a5L15oXp0w5L0smrLkzWTiyPx4uLrZAbTC+eqgfTtkcXFDlNdZFfqZNFFsppiadOqyk4hKBtMH58m0fiesfh9GmqJWnO0fgU3JMt6f98y/EiD8pEPFat5tJfqyY74nn4k+1F6m3xfY92xU+oobPSz9f3LpkIt131nfQ+//lL92twaP+MT0he2YFMEHYgE4QdyARhBzJB2IFMEHYgE4QdyERN57MXM3HseFhv/7f0PN70ArhTmi9cGdZHlsdj/KML0ruq5VT8O3O4Jx6HL2Y8PoW5JrvSY+XD8ZRwTbYUWepa8XjzRHv8FGoaTt9+03h8Lv/Rnni8Wa1FxuGb0vvFgpok6Vh8HgBvj3trKvKccEvf/8LX4n1q24L57OWcN97MVpjZT81sl5m9amZ3FK7vNrMtZra7cLmo2G0BqJ/ZvI0fl3SXu39K0q9Lus3MLpZ0t6St7r5K0tbC9wAaVNGwu/tBd3+x8PUJSbskLZO0XtLmwo9tlnR9lXoEUAEf6QM6MztX0mckbZe0xN0PSlO/ECQtTmyz0cz6zax/TPFx1ACqZ9ZhN7N5kn4o6U53T6+weBp33+Tufe7e16oiMz4AVM2swm5mrZoK+nfd/bHC1YfMrLdQ75UUn54VQF0VHXozM5P0kKRd7n7ftNKTkm6RdG/h8omqdFghE6/Hy+C2vB5vH+2oeAAJaAyzGWdfK+lmSa+Y2Y7CdV/WVMh/YGa3Stor6caqdAigIoqG3d2fk5Q6MiJ9JgoADYXDZYFMEHYgE4QdyARhBzJB2IFMEHYgE4QdyARhBzJB2IFMEHYgE4QdyARhBzJB2IFMEHYgE4QdyARhBzJB2IFMEHYgE4QdyARhBzJB2IFMEHYgE4QdyARhBzJB2IFMEHYgE4QdyARhBzJB2IFMFA27ma0ws5+a2S4ze9XM7ihcf4+Z7TezHYV/66rfLoBSzWZ99nFJd7n7i2Y2X9ILZralUPuGu3+teu0BqJTZrM9+UNLBwtcnzGyXpGXVbgxAZX2kv9nN7FxJn5G0vXDV7Wb2spk9bGaLEttsNLN+M+sf00h53QIo2azDbmbzJP1Q0p3uPijpfkkXSFqtqVf+r8+0nbtvcvc+d+9r1ZzyOwZQklmF3cxaNRX077r7Y5Lk7ofcfcLdJyU9KGlN9doEUK7ZfBpvkh6StMvd75t2fe+0H7tB0s7KtwegUmbzafxaSTdLesXMdhSu+7KkDWa2WpJL2iPpS1XoD0CFzObT+Ock2QylpyrfDoBq4Qg6IBOEHcgEYQcyQdiBTBB2IBOEHcgEYQcyQdiBTBB2IBOEHcgEYQcyQdiBTBB2IBOEHciEuXvt7szsXUn/M+2qHkmHa9bAR9OovTVqXxK9laqSvZ3j7mfOVKhp2D9052b97t5XtwYCjdpbo/Yl0VupatUbb+OBTBB2IBP1DvumOt9/pFF7a9S+JHorVU16q+vf7ABqp96v7ABqhLADmahL2M3sWjN73czeNLO769FDipntMbNXCstQ99e5l4fNbMDMdk67rtvMtpjZ7sLljGvs1am3hljGO1hmvK77rt7Ln9f8b3Yza5b0hqTPS9on6XlJG9z9tZo2kmBmeyT1uXvdD8Aws9+UNCTp2+7+6cJ1fy/pqLvfW/hFucjd/6JBertH0lC9l/EurFbUO32ZcUnXS/pD1XHfBX39vmqw3+rxyr5G0pvu/ra7j0r6nqT1deij4bn7s5KOnnb1ekmbC19v1tSTpeYSvTUEdz/o7i8Wvj4h6YNlxuu674K+aqIeYV8m6Z1p3+9TY6337pJ+YmYvmNnGejczgyXuflCaevJIWlznfk5XdBnvWjptmfGG2XelLH9ernqEfaalpBpp/G+tu18u6TpJtxXermJ2ZrWMd63MsMx4Qyh1+fNy1SPs+yStmPb9ckkH6tDHjNz9QOFyQNLjarylqA99sIJu4XKgzv38v0ZaxnumZcbVAPuunsuf1yPsz0taZWbnmVmbpJskPVmHPj7EzDoLH5zIzDolfUGNtxT1k5JuKXx9i6Qn6tjLL2mUZbxTy4yrzvuu7sufu3vN/0lap6lP5N+S9JV69JDo63xJLxX+vVrv3iQ9qqm3dWOaekd0q6QzJG2VtLtw2d1AvX1H0iuSXtZUsHrr1NtVmvrT8GVJOwr/1tV73wV91WS/cbgskAmOoAMyQdiBTBB2IBOEHcgEYQcyQdiBTBB2IBP/B3dg8QLkDMYGAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "i = 11\n",
    "plt.imshow(xb[i].view(28, 28))\n",
    "dsd['train'].features['label'].int2str(yb[i].item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d54e69a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "af653b8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import fastcore.all as fc\n",
    "import torch.optim as optim\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4f7aa326",
   "metadata": {},
   "outputs": [],
   "source": [
    "def_device = 'mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "32984ff4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'cpu'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def_device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b2018518",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Learner:\n",
    "    def __init__(self, model, dls, loss_func, lr, opt_func=optim.SGD): fc.store_attr()\n",
    "    def calc_stats(self):\n",
    "        acc = (self.preds.argmax(dim=1) == self.yb).float().sum()\n",
    "        n = len(self.xb)\n",
    "        self.accs.append(acc)\n",
    "        self.losses.append(self.loss*n)\n",
    "        self.ns.append(n)\n",
    "    \n",
    "    def one_batch(self):\n",
    "        self.xb,self.yb = self.batch\n",
    "        self.preds = self.model(self.xb)\n",
    "        self.loss = self.loss_func(self.preds, self.yb)\n",
    "        if self.model.training:\n",
    "            self.loss.backward()\n",
    "            self.opt.step()\n",
    "            self.opt.zero_grad()\n",
    "        with torch.no_grad():\n",
    "            self.calc_stats()  \n",
    "    \n",
    "    def one_epoch(self, train):\n",
    "        self.model.training = train\n",
    "        dl = self.dls.train if train else self.dls.valid\n",
    "        for self.num, self.batch in enumerate(dl): \n",
    "            self.one_batch()\n",
    "        n = sum(self.ns)\n",
    "        print(f'training:{self.model.training}, epoch:{self.epoch}, loss:{sum(self.losses).item()/n},\\\n",
    "              accuracy:{sum(self.accs).item()/n}')\n",
    "    \n",
    "    def fit(self, n_epochs):\n",
    "        self.accs,self.losses,self.ns = [],[],[]\n",
    "        self.model.to(def_device)\n",
    "        self.opt = self.opt_func(self.model.parameters(), lr=self.lr)\n",
    "        self.n_epochs = n_epochs\n",
    "        for self.epoch in range(self.n_epochs):\n",
    "            self.one_epoch(True)\n",
    "        with torch.no_grad():\n",
    "            self.one_epoch(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "00ef9519",
   "metadata": {},
   "outputs": [],
   "source": [
    "n,nh = 28*28,50\n",
    "model = nn.Sequential(*[nn.Linear(n, nh), nn.ReLU(nh), nn.Linear(nh, 10)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "eca3ed23",
   "metadata": {},
   "outputs": [],
   "source": [
    "learn = Learner(model, dls, loss_func=F.cross_entropy, lr=0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8370d6c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "training:True, epoch:0, loss:1.1884450520833334,              accuracy:0.5995166666666667\n",
      "training:True, epoch:1, loss:0.9456514973958333,              accuracy:0.674475\n",
      "training:False, epoch:1, loss:0.9268530048076923,              accuracy:0.6796\n"
     ]
    }
   ],
   "source": [
    "learn.fit(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "bd738d40",
   "metadata": {},
   "outputs": [],
   "source": [
    "from operator import attrgetter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "abc4b2e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_cbs(cbs, method_nm, learn=None):\n",
    "    for cb in sorted(cbs, key=attrgetter('order')):\n",
    "        method = getattr(cb, method_nm, None)\n",
    "        if method is not None:\n",
    "            method(learn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f40b8e7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CancelFitException(Exception):\n",
    "    pass\n",
    "class CancelBatchException(Exception):\n",
    "    pass\n",
    "class CancelEpochException(Exception):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "9f96abe0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Callback:\n",
    "    order = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "a7fa6dd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CompletionCB(Callback):\n",
    "    def before_fit(self, learn):\n",
    "        self.count = 0\n",
    "        \n",
    "    def after_batch(self, learn):\n",
    "        self.count += 1\n",
    "    \n",
    "    def after_fit(self, learn):\n",
    "        print(f'Completed {self.count} batches')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "e02d76f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "n,nh = 28*28,50\n",
    "def get_model():\n",
    "    return nn.Sequential(*[nn.Linear(n, nh), nn.ReLU(nh), nn.Linear(nh, 10)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "7c60c578",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Metric:\n",
    "    def __init__(self):\n",
    "        self.reset()\n",
    "    \n",
    "    def reset(self):\n",
    "        self.ns,self.vals = [],[]\n",
    "    \n",
    "    def add(self, inp, targ=None, n=1):\n",
    "        self.last = self.calc(inp, targ)\n",
    "        self.vals.append(self.last)\n",
    "        self.ns.append(n)\n",
    "    \n",
    "    def calc(self, inp, targ):\n",
    "        return inp\n",
    "\n",
    "    @property\n",
    "    def value(self):\n",
    "        ns = torch.tensor(self.ns)\n",
    "        vals = torch.tensor(self.vals)\n",
    "        return (vals*ns).sum()/ns.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "69aa0542",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Accuracy(Metric):\n",
    "    def calc(self, inp, targ):\n",
    "        return (inp== targ).float().mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "0b929dff",
   "metadata": {},
   "outputs": [],
   "source": [
    "acc = Accuracy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "7124c794",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.2667)"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "acc.add(torch.tensor([3, 1, 2, 0, 1, 2]), torch.tensor([0, 1, 1, 2, 1, 0]))\n",
    "acc.add(torch.tensor([1, 1, 2, 0, 5]), torch.tensor([0, 1, 1, 2, 1]))\n",
    "acc.value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "b52cbda8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.6176), 0.62)"
      ]
     },
     "execution_count": 82,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = Metric()\n",
    "loss.add(0.6, n=32)\n",
    "loss.add(0.9, n=2)\n",
    "loss.value, round((0.6*32+0.9*2)/(32+2), 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "178faad3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torcheval.metrics import MulticlassAccuracy, Mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "90cf6bd8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.2727)"
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metric = MulticlassAccuracy()\n",
    "metric.update(torch.tensor([3, 1, 2, 0, 1, 2]), torch.tensor([0, 1, 1, 2, 1, 0]))\n",
    "metric.update(torch.tensor([1, 1, 2, 0, 5]), torch.tensor([0, 1, 1, 2, 1]))\n",
    "metric.compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "60c20c8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections.abc import Mapping\n",
    "from copy import copy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "5ea27bc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_cpu(x):\n",
    "    if isinstance(x, Mapping): \n",
    "        return {k:to_cpu(v) for k, v in x.items()}\n",
    "    if isinstance(x, list):\n",
    "        return [to_cpu(ele) for ele in x]\n",
    "    if isinstance(x, tuple):\n",
    "        return tuple(to_cpu(list(x)))\n",
    "    res = x.detach().cpu()\n",
    "    return res.float() if res.dtype == torch.float16 else res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "18a04bcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MetricsCB(Callback):\n",
    "    def __init__(self, *ms, **metrics):\n",
    "        for o in ms: metrics[o.__name__] = o\n",
    "        self.metrics = metrics\n",
    "        self.all_metrics = copy(metrics)\n",
    "        self.all_metrics['loss'] = self.loss = Mean()\n",
    "    \n",
    "    def _log(self, d): print(d)\n",
    "    def before_fit(self, learn): learn.metrics = self\n",
    "    def before_epoch(self, learn): [o.reset() for o in self.all_metrics.values()]\n",
    "    \n",
    "    def after_epoch(self, learn):\n",
    "        log = {k:f'{v.compute()}' for k, v in self.all_metrics.items()}\n",
    "        log['epoch'] = learn.epoch\n",
    "        log['train'] = 'train' if learn.model.training else 'eval'\n",
    "        self._log(log)\n",
    "    \n",
    "    def after_batch(self, learn):\n",
    "        x,y,*_ = to_cpu(learn.batch)\n",
    "        for m in self.metrics.values(): m.update(to_cpu(learn.preds), y)\n",
    "        self.loss.update(to_cpu(learn.loss), weight=len(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "ec046af9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Learner:\n",
    "    def __init__(self, model, dls, loss_func, lr, cbs=[], opt_func=optim.SGD): fc.store_attr()\n",
    "    def calc_stats(self):\n",
    "        acc = (self.preds.argmax(dim=1) == self.yb).float().sum()\n",
    "        n = len(self.xb)\n",
    "        self.accs.append(acc)\n",
    "        self.losses.append(self.loss*n)\n",
    "        self.ns.append(n)\n",
    "    \n",
    "    def one_batch(self):\n",
    "        self.xb,self.yb = self.batch\n",
    "        self.preds = self.model(self.xb)\n",
    "        self.loss = self.loss_func(self.preds, self.yb)\n",
    "        if self.model.training:\n",
    "            self.loss.backward()\n",
    "            self.opt.step()\n",
    "            self.opt.zero_grad()\n",
    "        with torch.no_grad():\n",
    "            self.calc_stats()  \n",
    "    \n",
    "    def one_epoch(self, train):\n",
    "        self.model.training = train\n",
    "        dl = self.dls.train if train else self.dls.valid\n",
    "        try:\n",
    "            self.callback('before_epoch')\n",
    "            for self.num, self.batch in enumerate(dl): \n",
    "                try:\n",
    "                    self.callback('before_batch')\n",
    "                    self.one_batch()\n",
    "                    self.callback('after_batch')\n",
    "                except CancelBatchException:\n",
    "                    pass\n",
    "            self.callback('after_epoch')\n",
    "        except CancelEpochException:\n",
    "            pass\n",
    "       \n",
    "    def fit(self, n_epochs):\n",
    "        self.accs,self.losses,self.ns = [],[],[]\n",
    "        self.model.to(def_device)\n",
    "        self.opt = self.opt_func(self.model.parameters(), lr=self.lr)\n",
    "        self.n_epochs = n_epochs\n",
    "        try:\n",
    "            self.callback('before_fit')\n",
    "            for self.epoch in range(self.n_epochs):\n",
    "                self.one_epoch(True)\n",
    "            with torch.no_grad():\n",
    "                self.one_epoch(False)\n",
    "            self.callback('after_fit')\n",
    "        except CancelFitException:\n",
    "            pass\n",
    "    \n",
    "    def callback(self, method_nm):\n",
    "        run_cbs(self.cbs, method_nm, self)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "cddd9207",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = get_model()\n",
    "metric = MetricsCB(accuracy=MulticlassAccuracy())\n",
    "learn = Learner(model, dls, lr=0.2, loss_func=F.cross_entropy, cbs=[metric])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "eda93e06",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'accuracy': '0.6327499747276306', 'loss': '1.171434762064616', 'epoch': 0, 'train': 'train'}\n",
      "{'accuracy': '0.7446500062942505', 'loss': '0.7085846827189127', 'epoch': 1, 'train': 'train'}\n",
      "{'accuracy': '0.7824000120162964', 'loss': '0.6112953526814778', 'epoch': 2, 'train': 'train'}\n",
      "{'accuracy': '0.7728000283241272', 'loss': '0.6063536010742188', 'epoch': 2, 'train': 'eval'}\n"
     ]
    }
   ],
   "source": [
    "learn.fit(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "e17be3db",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MetricsCB(Callback):\n",
    "    def __init__(self, *ms, **metrics):\n",
    "        for m in ms:\n",
    "            metrics[o.__name__] = o\n",
    "        self.metrics = metrics\n",
    "        self.all_metrics = copy(metrics)\n",
    "        self.all_metrics['loss'] = self.loss = Mean()\n",
    "    \n",
    "    def before_fit(self, learn):\n",
    "        learn.metrics = self\n",
    "    \n",
    "    def before_epoch(self, learn):\n",
    "        [o.reset() for o in self.all_metrics.values()]\n",
    "    \n",
    "    def after_batch(self, learn):\n",
    "        x,y,*_ = learn.batch\n",
    "        for m in self.metrics.values():\n",
    "            m.update(to_cpu(learn.preds), y)\n",
    "        self.loss.update(to_cpu(learn.loss), weight=len(x))\n",
    "    \n",
    "    def _log(self, d):\n",
    "        print(d)\n",
    "    \n",
    "    def after_epoch(self, learn):\n",
    "        log = {k:f'{v.compute():.3f}' for k, v in self.all_metrics.items()}\n",
    "        log['epoch'] = learn.epoch\n",
    "        log['train'] = 'train' if learn.model.training else 'eval'\n",
    "        self._log(log)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "d77abaf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = get_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "3c8d9a11",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = MetricsCB(accuracy=MulticlassAccuracy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "745ade9c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'accuracy': <torcheval.metrics.classification.accuracy.MulticlassAccuracy at 0x1794e9d50>,\n",
       " 'loss': <torcheval.metrics.aggregation.mean.Mean at 0x177385910>}"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metrics.all_metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "8ab200a3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'accuracy': '0.606', 'loss': '1.176', 'epoch': 0, 'train': 'train'}\n",
      "{'accuracy': '0.695', 'loss': '0.793', 'epoch': 0, 'train': 'eval'}\n"
     ]
    }
   ],
   "source": [
    "learn = Learner(model, dls, F.cross_entropy, lr=0.2, cbs=[metrics])\n",
    "learn.fit(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "c4d18755",
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_device(x, device=def_device):\n",
    "    if isinstance(x, torch.Tensor): return x.to(device)\n",
    "    if isinstance(x, Mapping): return {k:v.to(device) for k,v in x.items()}\n",
    "    return type(x)(to_device(o, device) for o in x)\n",
    "\n",
    "class DeviceCB(Callback):\n",
    "    def __init(self, device=def_device): fc.store_attr()\n",
    "    def before_fit(self, learn): \n",
    "        if hasattr(learn.model, 'to'):learn.model.to(self.device)\n",
    "    def before_batch(self, learn): learn.batch = to_device(learn.batch, device=self.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "42872ff2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_device(x, device=def_device):\n",
    "    if isinstance(x, torch.Tensor):\n",
    "        return x.to(device)\n",
    "    if isinstance(x, Mapping):\n",
    "        return {k:v.to(device) for k, v in x.items()}\n",
    "    return type(x)(to_device(o, device) for o in x)\n",
    "\n",
    "class DeviceCB(Callback):\n",
    "    def __init__(self, device=def_device):\n",
    "        fc.store_attr()\n",
    "    \n",
    "    def before_fit(self, learn):\n",
    "        if hasattr(learn.model, 'to'):\n",
    "            learn.model.to(self.device)\n",
    "    \n",
    "    def before_batch(self, learn):\n",
    "        learn.batch = to_device(learn.batch, device=self.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "c56a7149",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'accuracy': '0.746', 'loss': '0.710', 'epoch': 0, 'train': 'train'}\n",
      "{'accuracy': '0.735', 'loss': '0.707', 'epoch': 0, 'train': 'eval'}\n"
     ]
    }
   ],
   "source": [
    "learn = Learner(model, dls, F.cross_entropy, lr=0.2, cbs=[DeviceCB(), metrics])\n",
    "learn.fit(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "id": "3d5dc967",
   "metadata": {},
   "outputs": [],
   "source": [
    "from contextlib import contextmanager\n",
    "from functools import partial\n",
    "from fastprogress import progress_bar,master_bar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "893e280f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Learner:\n",
    "    def __init__(self, modele, dls=(0,), lr=0.1, loss_func=F.mse_loss, cbs=None, opt_func=optim.SGD):\n",
    "        cbs = fc.L(cbs)\n",
    "        fc.store_attr()\n",
    "    \n",
    "    @contextmanager\n",
    "    def cb_ctx(self, nm):\n",
    "        try:\n",
    "            self.callback(f'before_{nm}')\n",
    "            yield\n",
    "            self.callback(f'after_{nm}')\n",
    "        except globals()[f'Cancel{nm.title()}Exception']:\n",
    "            pass\n",
    "        finally:\n",
    "            self.callback(f'cleanup_{nm}')\n",
    "    \n",
    "    def one_epoch(self, train):\n",
    "        self.model.train(train)\n",
    "        self.dl = self.dls.train if train else self.dls.valid\n",
    "        with self.cb_ctx('epoch'):\n",
    "            for self.iter,self.epoch in enumerate(self.dl):\n",
    "                with self.cb_ctx('batch'):\n",
    "                    self.predict()\n",
    "                    self.get_loss()\n",
    "                    if self.training:\n",
    "                        self.backward()\n",
    "                        self.step()\n",
    "                        self.zero_grad()\n",
    "    \n",
    "    def fit(self, model, train=True, valid=True, n_epochs=1, cbs=None, lr=None):\n",
    "        cbs = fc.L(cbs)\n",
    "        for cb in cbs:\n",
    "            self.cbs.append(cb)\n",
    "        try:\n",
    "            self.n_epochs = n_epochs\n",
    "            self.epochs = range(n_epochs)\n",
    "            self.opt = self.opt_func(model.parameters(), self.lr if lr is None else lr)\n",
    "            with self.cb_ctx('fit'):\n",
    "                for self.epoch in self.epochs:\n",
    "                    if train:\n",
    "                        self.one_epoch(True)\n",
    "                    if valid:\n",
    "                        with torch.no_grad():\n",
    "                            self.one_epoch(False)\n",
    "        finally:\n",
    "            for cb in cbs:\n",
    "                self.cbs.remove(cb)\n",
    "    \n",
    "    def __getattr__(self, name):\n",
    "        if name in ('predict','get_loss','backward','step','zero_grad'):\n",
    "            return partial(self.callback, name)\n",
    "        raise AttributeError(name)\n",
    "        \n",
    "    def callback(self, method_nm):\n",
    "        run_cbs(self.cbs, method_nm, self)\n",
    "    \n",
    "    @property\n",
    "    def training(self):\n",
    "        return self.model.training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "id": "a93f2b1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrainCB(Callback):\n",
    "    def __init__(self, n_inp=1):\n",
    "        self.n_inp = n_inp\n",
    "    \n",
    "    def predict(self, learn):\n",
    "        learn.preds = learn.model(*learn.batch[:self.n_inp])\n",
    "    \n",
    "    def get_loss(self, learn):\n",
    "        learn.loss = learn.loss_func(learn.preds, *learn.batch[self.n_inp:])\n",
    "    \n",
    "    def backward(self, learn):\n",
    "        learn.loss.backward()\n",
    "    \n",
    "    def step(self, learn):\n",
    "        learn.opt.step()\n",
    "    \n",
    "    def zero_grad(self, learn):\n",
    "        learn.opt.zero_grad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "id": "5496c444",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ProgressCB(Callback):\n",
    "    order = MetricsCB.order+1\n",
    "    def __init__(self, plot=False):\n",
    "        self.plot = plot\n",
    "    \n",
    "    def before_fit(self, learn):\n",
    "        learn.epochs = self.mbar = master_bar(learn.epochs)\n",
    "        self.first = True\n",
    "        if hasattr(learn, 'metrics'):\n",
    "            learn.metrics._log = self._log\n",
    "        self.losses = []\n",
    "        self.val_losses = []\n",
    "    \n",
    "    def _log(self, d):\n",
    "        if self.first:\n",
    "            self.mbar.write(list(d), table=True)\n",
    "            self.first =False\n",
    "        self.mbar.write(list(d.values()), table=True)\n",
    "    \n",
    "    def before_epoch(self, learn):\n",
    "        learn.dl = progress_bar(learn.dl, leave=False, parent=self.mbar)\n",
    "    \n",
    "    def after_batch(self, learn):\n",
    "        learn.dl.comment = f'{learn.loss:.3f}'\n",
    "        if self.plot and hasattr(learn, 'metrics') and learn.training:\n",
    "            self.losses.append(learn.loss.item())\n",
    "            if self.val_losses:\n",
    "                self.mbar.update_graph()\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "id": "8a9454c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Learner():\n",
    "    def __init__(self, model, dls=(0,), loss_func=F.mse_loss, lr=0.1, cbs=None, opt_func=optim.SGD):\n",
    "        cbs = fc.L(cbs)\n",
    "        fc.store_attr()\n",
    "\n",
    "    @contextmanager\n",
    "    def cb_ctx(self, nm):\n",
    "        try:\n",
    "            self.callback(f'before_{nm}')\n",
    "            yield\n",
    "            self.callback(f'after_{nm}')\n",
    "        except globals()[f'Cancel{nm.title()}Exception']: pass\n",
    "        finally: self.callback(f'cleanup_{nm}')\n",
    "                \n",
    "    def one_epoch(self, train):\n",
    "        self.model.train(train)\n",
    "        self.dl = self.dls.train if train else self.dls.valid\n",
    "        with self.cb_ctx('epoch'):\n",
    "            for self.iter,self.batch in enumerate(self.dl):\n",
    "                with self.cb_ctx('batch'):\n",
    "                    self.predict()\n",
    "                    self.get_loss()\n",
    "                    if self.training:\n",
    "                        self.backward()\n",
    "                        self.step()\n",
    "                        self.zero_grad()\n",
    "    \n",
    "    def fit(self, n_epochs=1, train=True, valid=True, cbs=None, lr=None):\n",
    "        cbs = fc.L(cbs)\n",
    "        # `add_cb` and `rm_cb` were added in lesson 18\n",
    "        for cb in cbs: self.cbs.append(cb)\n",
    "        try:\n",
    "            self.n_epochs = n_epochs\n",
    "            self.epochs = range(n_epochs)\n",
    "            self.opt = self.opt_func(self.model.parameters(), self.lr if lr is None else lr)\n",
    "            with self.cb_ctx('fit'):\n",
    "                for self.epoch in self.epochs:\n",
    "                    if train: self.one_epoch(True)\n",
    "                    if valid: torch.no_grad()(self.one_epoch)(False)\n",
    "        finally:\n",
    "            for cb in cbs: self.cbs.remove(cb)\n",
    "\n",
    "    def __getattr__(self, name):\n",
    "        if name in ('predict','get_loss','backward','step','zero_grad'): return partial(self.callback, name)\n",
    "        raise AttributeError(name)\n",
    "\n",
    "    def callback(self, method_nm): run_cbs(self.cbs, method_nm, self)\n",
    "    \n",
    "    @property\n",
    "    def training(self): return self.model.training\n",
    "     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "id": "282c76c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrainCB(Callback):\n",
    "    def __init__(self, n_inp=1): self.n_inp = n_inp\n",
    "    def predict(self, learn): learn.preds = learn.model(*learn.batch[:self.n_inp])\n",
    "    def get_loss(self, learn): learn.loss = learn.loss_func(learn.preds, *learn.batch[self.n_inp:])\n",
    "    def backward(self, learn): learn.loss.backward()\n",
    "    def step(self, learn): learn.opt.step()\n",
    "    def zero_grad(self, learn): learn.opt.zero_grad()\n",
    "     \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "id": "9cb983b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class ProgressCB(Callback):\n",
    "    order = MetricsCB.order+1\n",
    "    def __init__(self, plot=False): self.plot = plot\n",
    "    def before_fit(self, learn):\n",
    "        learn.epochs = self.mbar = master_bar(learn.epochs)\n",
    "        self.first = True\n",
    "        if hasattr(learn, 'metrics'): learn.metrics._log = self._log\n",
    "        self.losses = []\n",
    "        self.val_losses = []\n",
    "\n",
    "    def _log(self, d):\n",
    "        if self.first:\n",
    "            self.mbar.write(list(d), table=True)\n",
    "            self.first = False\n",
    "        self.mbar.write(list(d.values()), table=True)\n",
    "\n",
    "    def before_epoch(self, learn): learn.dl = progress_bar(learn.dl, leave=False, parent=self.mbar)\n",
    "    \n",
    "    def after_batch(self, learn):\n",
    "        learn.dl.comment = f'{learn.loss:.3f}'\n",
    "        if self.plot and hasattr(learn, 'metrics') and learn.training:\n",
    "            self.losses.append(learn.loss.item())\n",
    "            if self.val_losses:\n",
    "                self.mbar.update_graph([[fc.L.range(self.losses), self.losses],\\\n",
    "                    [fc.L.range(learn.epoch).map(lambda x: (x+1)*len(learn.dls.train)), self.val_losses]])\n",
    "   \n",
    "    def ater_epoch(self, learn):\n",
    "        if not learn.training:\n",
    "            self.val_losses.append(learn.metrics.all_metrics['loss'].compute())\n",
    "            self.mbar.update_graph([[fc.L.range(self.losses), self.losses],\\\n",
    "                    [fc.L.range(learn.epoch+1).map(lambda x: (x+1)*len(learn.dls.train)), self.val_losses]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "id": "a2700993",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = get_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "id": "06a23e53",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = MetricsCB(accuracy=MulticlassAccuracy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "id": "e64e6038",
   "metadata": {},
   "outputs": [],
   "source": [
    "cbs = [TrainCB(), DeviceCB(), metrics, ProgressCB(plot=True)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "id": "861a4bff",
   "metadata": {},
   "outputs": [],
   "source": [
    "learn = Learner(model, dls, loss_func=F.cross_entropy, cbs=cbs, lr=0.2, )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "id": "2c0e01d5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "<style>\n",
       "    /* Turns off some styling */\n",
       "    progress {\n",
       "        /* gets rid of default border in Firefox and Opera. */\n",
       "        border: none;\n",
       "        /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "        background-size: auto;\n",
       "    }\n",
       "    progress:not([value]), progress:not([value])::-webkit-progress-bar {\n",
       "        background: repeating-linear-gradient(45deg, #7e7e7e, #7e7e7e 10px, #5c5c5c 10px, #5c5c5c 20px);\n",
       "    }\n",
       "    .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "        background: #F44336;\n",
       "    }\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>accuracy</th>\n",
       "      <th>loss</th>\n",
       "      <th>epoch</th>\n",
       "      <th>train</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0.779</td>\n",
       "      <td>0.616</td>\n",
       "      <td>0</td>\n",
       "      <td>train</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>0.764</td>\n",
       "      <td>0.679</td>\n",
       "      <td>0</td>\n",
       "      <td>eval</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>0.800</td>\n",
       "      <td>0.567</td>\n",
       "      <td>1</td>\n",
       "      <td>train</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>0.796</td>\n",
       "      <td>0.564</td>\n",
       "      <td>1</td>\n",
       "      <td>eval</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "learn.fit(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "id": "e6ecf605",
   "metadata": {},
   "outputs": [],
   "source": [
    "class with_cbs:\n",
    "    def __init__(self, nm):\n",
    "        self.nm = nm\n",
    "    \n",
    "    def __call__(self, f):\n",
    "        def _f(o, *args, **kwargs):\n",
    "            try:\n",
    "                o.callback(f'before_{self.nm}')\n",
    "                f(o, *args, **kwargs)\n",
    "                o.callback(f'after_{self.nm}')\n",
    "            except globals()[f'Cacnel{self.nm.title()}Exception']:\n",
    "                pass\n",
    "            finally:\n",
    "                o.callback(f'cleanup_{self.nm}')\n",
    "        return _f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "id": "0c138106",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Learner():\n",
    "    def __init__(self, model, dls=(0,), loss_func=F.mse_loss, lr=0.1, cbs=None, opt_func=optim.SGD):\n",
    "        cbs = fc.L(cbs)\n",
    "        fc.store_attr()\n",
    "\n",
    "    @contextmanager\n",
    "    def cb_ctx(self, nm):\n",
    "        try:\n",
    "            self.callback(f'before_{nm}')\n",
    "            yield\n",
    "            self.callback(f'after_{nm}')\n",
    "        except globals()[f'Cancel{nm.title()}Exception']: pass\n",
    "        finally: self.callback(f'cleanup_{nm}')\n",
    "                \n",
    "    def one_epoch(self, train):\n",
    "        self.model.train(train)\n",
    "        self.dl = self.dls.train if train else self.dls.valid\n",
    "        with self.cb_ctx('epoch'):\n",
    "            for self.iter,self.batch in enumerate(self.dl):\n",
    "                with self.cb_ctx('batch'):\n",
    "                    self.predict()\n",
    "                    self.get_loss()\n",
    "                    if self.training:\n",
    "                        self.backward()\n",
    "                        self.step()\n",
    "                        self.zero_grad()\n",
    "    \n",
    "    def fit(self, n_epochs=1, train=True, valid=True, cbs=None, lr=None):\n",
    "        cbs = fc.L(cbs)\n",
    "        # `add_cb` and `rm_cb` were added in lesson 18\n",
    "        for cb in cbs: self.cbs.append(cb)\n",
    "        try:\n",
    "            self.n_epochs = n_epochs\n",
    "            self.epochs = range(n_epochs)\n",
    "            self.opt = self.opt_func(self.model.parameters(), self.lr if lr is None else lr)\n",
    "            with self.cb_ctx('fit'):\n",
    "                for self.epoch in self.epochs:\n",
    "                    if train: self.one_epoch(True)\n",
    "                    if valid: torch.no_grad()(self.one_epoch)(False)\n",
    "        finally:\n",
    "            for cb in cbs: self.cbs.remove(cb)\n",
    "\n",
    "    def __getattr__(self, name):\n",
    "        if name in ('predict','get_loss','backward','step','zero_grad'): return partial(self.callback, name)\n",
    "        raise AttributeError(name)\n",
    "\n",
    "    def callback(self, method_nm): run_cbs(self.cbs, method_nm, self)\n",
    "    \n",
    "    @property\n",
    "    def training(self): return self.model.training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "id": "dd451a01",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrainCB(Callback):\n",
    "    def __init__(self, n_inp=1): self.n_inp = n_inp\n",
    "    def predict(self, learn): learn.preds = learn.model(*learn.batch[:self.n_inp])\n",
    "    def get_loss(self, learn): learn.loss = learn.loss_func(learn.preds, *learn.batch[self.n_inp:])\n",
    "    def backward(self, learn): learn.loss.backward()\n",
    "    def step(self, learn): learn.opt.step()\n",
    "    def zero_grad(self, learn): learn.opt.zero_grad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "id": "8316a0fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = get_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "id": "10277c76",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "<style>\n",
       "    /* Turns off some styling */\n",
       "    progress {\n",
       "        /* gets rid of default border in Firefox and Opera. */\n",
       "        border: none;\n",
       "        /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "        background-size: auto;\n",
       "    }\n",
       "    progress:not([value]), progress:not([value])::-webkit-progress-bar {\n",
       "        background: repeating-linear-gradient(45deg, #7e7e7e, #7e7e7e 10px, #5c5c5c 10px, #5c5c5c 20px);\n",
       "    }\n",
       "    .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "        background: #F44336;\n",
       "    }\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>accuracy</th>\n",
       "      <th>loss</th>\n",
       "      <th>epoch</th>\n",
       "      <th>train</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0.618</td>\n",
       "      <td>1.185</td>\n",
       "      <td>0</td>\n",
       "      <td>train</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>0.684</td>\n",
       "      <td>0.809</td>\n",
       "      <td>0</td>\n",
       "      <td>eval</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>0.735</td>\n",
       "      <td>0.721</td>\n",
       "      <td>1</td>\n",
       "      <td>train</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>0.736</td>\n",
       "      <td>0.711</td>\n",
       "      <td>1</td>\n",
       "      <td>eval</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "metrics = MetricsCB(accuracy=MulticlassAccuracy())\n",
    "cbs = [TrainCB(), DeviceCB(), metrics, ProgressCB(plot=True)]\n",
    "learn = Learner(model, dls, F.cross_entropy, lr=0.2, cbs=cbs)\n",
    "learn.fit(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c852b59",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python_main",
   "language": "python",
   "name": "python_main"
  },
  "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.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}