lec_4_v2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 105,
   "id": "249fd15c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch, pickle, matplotlib as mpl, gzip\n",
    "from pathlib import Path\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "mpl.rcParams['image.cmap'] = 'gray'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ec268e2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = None\n",
    "with gzip.open('data/mnist.pkl.gz', 'rb') as f:\n",
    "    data = pickle.load(f, encoding='latin')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f38a580d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((50000, 784), (50000,), (10000, 784), (10000,))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "((x_train, y_train), (x_val, y_val), _) = data\n",
    "x_train.shape, y_train.shape, x_val.shape, y_val.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "758f345a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([50000, 784]),\n",
       " torch.Size([50000]),\n",
       " torch.Size([10000, 784]),\n",
       " torch.Size([10000]))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(x_train, y_train, x_val, y_val) = map(torch.tensor, (x_train, y_train, x_val, y_val))\n",
    "x_train.shape, y_train.shape, x_val.shape, y_val.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a3d22db1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x139bacb10>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD4CAYAAAAq5pAIAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAAsTAAALEwEAmpwYAAAN80lEQVR4nO3df6hcdXrH8c+ncf3DrBpTMYasNhuRWBWbLRqLSl2RrD9QNOqWDVgsBrN/GHChhEr6xyolEuqP0qAsuYu6sWyzLqgYZVkVo6ZFCF5j1JjU1YrdjV6SSozG+KtJnv5xT+Su3vnOzcyZOZP7vF9wmZnzzJnzcLife87Md879OiIEYPL7k6YbANAfhB1IgrADSRB2IAnCDiRxRD83ZpuP/oEeiwiPt7yrI7vtS22/aftt27d281oAesudjrPbniLpd5IWSNou6SVJiyJia2EdjuxAj/XiyD5f0tsR8U5EfCnpV5Ku6uL1APRQN2GfJekPYx5vr5b9EdtLbA/bHu5iWwC61M0HdOOdKnzjND0ihiQNSZzGA03q5si+XdJJYx5/R9L73bUDoFe6CftLkk61/V3bR0r6kaR19bQFoG4dn8ZHxD7bSyU9JWmKpAci4o3aOgNQq46H3jraGO/ZgZ7ryZdqABw+CDuQBGEHkiDsQBKEHUiCsANJEHYgCcIOJEHYgSQIO5AEYQeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgrADSRB2IAnCDiRB2IEkCDuQBGEHkiDsQBKEHUii4ymbcXiYMmVKsX7sscf2dPtLly5tWTvqqKOK686dO7dYv/nmm4v1u+66q2Vt0aJFxXU///zzYn3lypXF+u23316sN6GrsNt+V9IeSfsl7YuIs+toCkD96jiyXxQRH9TwOgB6iPfsQBLdhj0kPW37ZdtLxnuC7SW2h20Pd7ktAF3o9jT+/Ih43/YJkp6x/V8RsWHsEyJiSNKQJNmOLrcHoENdHdkj4v3qdqekxyTNr6MpAPXrOOy2p9o++uB9ST+QtKWuxgDUq5vT+BmSHrN98HX+PSJ+W0tXk8zJJ59crB955JHF+nnnnVesX3DBBS1r06ZNK6577bXXFutN2r59e7G+atWqYn3hwoUta3v27Cmu++qrrxbrL7zwQrE+iDoOe0S8I+kvauwFQA8x9AYkQdiBJAg7kARhB5Ig7EASjujfl9om6zfo5s2bV6yvX7++WO/1ZaaD6sCBA8X6jTfeWKx/8sknHW97ZGSkWP/www+L9TfffLPjbfdaRHi85RzZgSQIO5AEYQeSIOxAEoQdSIKwA0kQdiAJxtlrMH369GJ948aNxfqcOXPqbKdW7XrfvXt3sX7RRRe1rH355ZfFdbN+/6BbjLMDyRF2IAnCDiRB2IEkCDuQBGEHkiDsQBJM2VyDXbt2FevLli0r1q+44opi/ZVXXinW2/1L5ZLNmzcX6wsWLCjW9+7dW6yfccYZLWu33HJLcV3UiyM7kARhB5Ig7EAShB1IgrADSRB2IAnCDiTB9ewD4JhjjinW200vvHr16pa1xYsXF9e9/vrri/W1a9cW6xg8HV/PbvsB2zttbxmzbLrtZ2y/Vd0eV2ezAOo3kdP4X0i69GvLbpX0bEScKunZ6jGAAdY27BGxQdLXvw96laQ11f01kq6uty0Adev0u/EzImJEkiJixPYJrZ5oe4mkJR1uB0BNen4hTEQMSRqS+IAOaFKnQ287bM+UpOp2Z30tAeiFTsO+TtIN1f0bJD1eTzsAeqXtabzttZK+L+l429sl/VTSSkm/tr1Y0u8l/bCXTU52H3/8cVfrf/TRRx2ve9NNNxXrDz/8cLHebo51DI62YY+IRS1KF9fcC4Ae4uuyQBKEHUiCsANJEHYgCcIOJMElrpPA1KlTW9aeeOKJ4roXXnhhsX7ZZZcV608//XSxjv5jymYgOcIOJEHYgSQIO5AEYQeSIOxAEoQdSIJx9knulFNOKdY3bdpUrO/evbtYf+6554r14eHhlrX77ruvuG4/fzcnE8bZgeQIO5AEYQeSIOxAEoQdSIKwA0kQdiAJxtmTW7hwYbH+4IMPFutHH310x9tevnx5sf7QQw8V6yMjIx1vezJjnB1IjrADSRB2IAnCDiRB2IEkCDuQBGEHkmCcHUVnnnlmsX7PPfcU6xdf3Plkv6tXry7WV6xYUay/9957HW/7cNbxOLvtB2zvtL1lzLLbbL9ne3P1c3mdzQKo30RO438h6dJxlv9LRMyrfn5Tb1sA6tY27BGxQdKuPvQCoIe6+YBuqe3XqtP841o9yfYS28O2W/8zMgA912nYfybpFEnzJI1IurvVEyNiKCLOjoizO9wWgBp0FPaI2BER+yPigKSfS5pfb1sA6tZR2G3PHPNwoaQtrZ4LYDC0HWe3vVbS9yUdL2mHpJ9Wj+dJCknvSvpxRLS9uJhx9sln2rRpxfqVV17ZstbuWnl73OHir6xfv75YX7BgQbE+WbUaZz9iAisuGmfx/V13BKCv+LoskARhB5Ig7EAShB1IgrADSXCJKxrzxRdfFOtHHFEeLNq3b1+xfskll7SsPf/888V1D2f8K2kgOcIOJEHYgSQIO5AEYQeSIOxAEoQdSKLtVW/I7ayzzirWr7vuumL9nHPOaVlrN47eztatW4v1DRs2dPX6kw1HdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgnH2SW7u3LnF+tKlS4v1a665plg/8cQTD7mnidq/f3+xPjJS/u/lBw4cqLOdwx5HdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgnH2w0C7sexFi8abaHdUu3H02bNnd9JSLYaHh4v1FStWFOvr1q2rs51Jr+2R3fZJtp+zvc32G7ZvqZZPt/2M7beq2+N63y6ATk3kNH6fpL+PiD+X9FeSbrZ9uqRbJT0bEadKerZ6DGBAtQ17RIxExKbq/h5J2yTNknSVpDXV09ZIurpHPQKowSG9Z7c9W9L3JG2UNCMiRqTRPwi2T2ixzhJJS7rsE0CXJhx229+W9Iikn0TEx/a4c8d9Q0QMSRqqXoOJHYGGTGjozfa3NBr0X0bEo9XiHbZnVvWZknb2pkUAdWh7ZPfoIfx+Sdsi4p4xpXWSbpC0srp9vCcdTgIzZswo1k8//fRi/d577y3WTzvttEPuqS4bN24s1u+8886WtccfL//KcIlqvSZyGn++pL+V9LrtzdWy5RoN+a9tL5b0e0k/7EmHAGrRNuwR8Z+SWr1Bv7jedgD0Cl+XBZIg7EAShB1IgrADSRB2IAkucZ2g6dOnt6ytXr26uO68efOK9Tlz5nTSUi1efPHFYv3uu+8u1p966qli/bPPPjvkntAbHNmBJAg7kARhB5Ig7EAShB1IgrADSRB2IIk04+znnntusb5s2bJiff78+S1rs2bN6qinunz66acta6tWrSque8cddxTre/fu7agnDB6O7EAShB1IgrADSRB2IAnCDiRB2IEkCDuQRJpx9oULF3ZV78bWrVuL9SeffLJY37dvX7FeuuZ89+7dxXWRB0d2IAnCDiRB2IEkCDuQBGEHkiDsQBKEHUjCEVF+gn2SpIcknSjpgKShiPhX27dJuknS/1ZPXR4Rv2nzWuWNAehaRIw76/JEwj5T0syI2GT7aEkvS7pa0t9I+iQi7ppoE4Qd6L1WYZ/I/Owjkkaq+3tsb5PU7L9mAXDIDuk9u+3Zkr4naWO1aKnt12w/YPu4FusssT1se7i7VgF0o+1p/FdPtL8t6QVJKyLiUdszJH0gKST9k0ZP9W9s8xqcxgM91vF7dkmy/S1JT0p6KiLuGac+W9KTEXFmm9ch7ECPtQp729N425Z0v6RtY4NefXB30EJJW7ptEkDvTOTT+Ask/Yek1zU69CZJyyUtkjRPo6fx70r6cfVhXum1OLIDPdbVaXxdCDvQex2fxgOYHAg7kARhB5Ig7EAShB1IgrADSRB2IAnCDiRB2IEkCDuQBGEHkiDsQBKEHUiCsANJ9HvK5g8k/c+Yx8dXywbRoPY2qH1J9NapOnv7s1aFvl7P/o2N28MRcXZjDRQMam+D2pdEb53qV2+cxgNJEHYgiabDPtTw9ksGtbdB7Uuit071pbdG37MD6J+mj+wA+oSwA0k0Enbbl9p+0/bbtm9toodWbL9r+3Xbm5uen66aQ2+n7S1jlk23/Yztt6rbcefYa6i322y/V+27zbYvb6i3k2w/Z3ub7Tds31Itb3TfFfrqy37r+3t221Mk/U7SAknbJb0kaVFEbO1rIy3YflfS2RHR+BcwbP+1pE8kPXRwai3b/yxpV0SsrP5QHhcR/zAgvd2mQ5zGu0e9tZpm/O/U4L6rc/rzTjRxZJ8v6e2IeCcivpT0K0lXNdDHwIuIDZJ2fW3xVZLWVPfXaPSXpe9a9DYQImIkIjZV9/dIOjjNeKP7rtBXXzQR9lmS/jDm8XYN1nzvIelp2y/bXtJ0M+OYcXCarer2hIb7+bq203j309emGR+YfdfJ9OfdaiLs401NM0jjf+dHxF9KukzSzdXpKibmZ5JO0egcgCOS7m6ymWqa8Uck/SQiPm6yl7HG6asv+62JsG+XdNKYx9+R9H4DfYwrIt6vbndKekyjbzsGyY6DM+hWtzsb7ucrEbEjIvZHxAFJP1eD+66aZvwRSb+MiEerxY3vu/H66td+ayLsL0k61fZ3bR8p6UeS1jXQxzfYnlp9cCLbUyX9QIM3FfU6STdU92+Q9HiDvfyRQZnGu9U042p43zU+/XlE9P1H0uUa/UT+vyX9YxM9tOhrjqRXq583mu5N0lqNntb9n0bPiBZL+lNJz0p6q7qdPkC9/ZtGp/Z+TaPBmtlQbxdo9K3ha5I2Vz+XN73vCn31Zb/xdVkgCb5BByRB2IEkCDuQBGEHkiDsQBKEHUiCsANJ/D+f1mbt6t55/AAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(x_train[0].view(28, 28))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "16cee6ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self, nin, nh, nout):\n",
    "        super().__init__()\n",
    "        self.layers = [nn.Linear(nin, nh), nn.ReLU(nh), nn.Linear(nh, nout)]\n",
    "    \n",
    "    def __call__(self, x):\n",
    "        for l in self.layers:\n",
    "            x = l(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "97c55568",
   "metadata": {},
   "outputs": [],
   "source": [
    "m, n = x_train.shape\n",
    "c = y_train.max() + 1\n",
    "nh = 50\n",
    "nout = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "462fbc99",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model(n, nh, nout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "b9455fde",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = model(x_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "0ae35781",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([50000, 10])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "b693b124",
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_softmax(x):\n",
    "    return (x.exp()/x.exp().sum(1, keepdim=True)).log()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "43dafa99",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-2.2446, -2.4244, -2.4216,  ..., -2.2878, -2.4068, -2.3208],\n",
       "        [-2.3862, -2.3740, -2.3016,  ..., -2.2498, -2.3517, -2.4064],\n",
       "        [-2.3970, -2.4087, -2.2812,  ..., -2.3484, -2.4407, -2.3838],\n",
       "        ...,\n",
       "        [-2.3303, -2.4678, -2.2690,  ..., -2.4128, -2.2866, -2.4098],\n",
       "        [-2.2764, -2.4901, -2.2623,  ..., -2.3608, -2.3901, -2.3783],\n",
       "        [-2.3099, -2.3961, -2.2686,  ..., -2.3862, -2.3096, -2.4644]],\n",
       "       grad_fn=<LogBackward0>)"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_softmax(pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "ee2e819d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_softmax(x):\n",
    "    return x - (x.exp().sum(1, keepdim=True)).log()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "4113a758",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-2.2446, -2.4244, -2.4216,  ..., -2.2878, -2.4068, -2.3208],\n",
       "        [-2.3862, -2.3740, -2.3016,  ..., -2.2498, -2.3517, -2.4064],\n",
       "        [-2.3970, -2.4087, -2.2812,  ..., -2.3484, -2.4407, -2.3838],\n",
       "        ...,\n",
       "        [-2.3303, -2.4678, -2.2690,  ..., -2.4128, -2.2866, -2.4098],\n",
       "        [-2.2764, -2.4901, -2.2623,  ..., -2.3608, -2.3901, -2.3783],\n",
       "        [-2.3099, -2.3961, -2.2686,  ..., -2.3862, -2.3096, -2.4644]],\n",
       "       grad_fn=<SubBackward0>)"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_softmax(pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "859887d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def logexpsum(x):\n",
    "    m = x.max(1, keepdim=True)[0]\n",
    "    return m + (x - m).exp().sum(1, keepdim=True).log()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "ba12784d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_softmax(x):\n",
    "    return x - logexpsum(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "aa3c98f3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-2.2446, -2.4244, -2.4216,  ..., -2.2878, -2.4068, -2.3208],\n",
       "        [-2.3862, -2.3740, -2.3016,  ..., -2.2498, -2.3517, -2.4064],\n",
       "        [-2.3970, -2.4087, -2.2812,  ..., -2.3484, -2.4407, -2.3838],\n",
       "        ...,\n",
       "        [-2.3303, -2.4678, -2.2690,  ..., -2.4128, -2.2866, -2.4098],\n",
       "        [-2.2764, -2.4901, -2.2623,  ..., -2.3608, -2.3901, -2.3783],\n",
       "        [-2.3099, -2.3961, -2.2686,  ..., -2.3862, -2.3096, -2.4644]],\n",
       "       grad_fn=<SubBackward0>)"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_softmax(pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "18c6ad7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_softmax(x):\n",
    "    return x - x.logsumexp(1, keepdim=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "dd276622",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-2.2446, -2.4244, -2.4216,  ..., -2.2878, -2.4068, -2.3208],\n",
       "        [-2.3862, -2.3740, -2.3016,  ..., -2.2498, -2.3517, -2.4064],\n",
       "        [-2.3970, -2.4087, -2.2812,  ..., -2.3484, -2.4407, -2.3838],\n",
       "        ...,\n",
       "        [-2.3303, -2.4678, -2.2690,  ..., -2.4128, -2.2866, -2.4098],\n",
       "        [-2.2764, -2.4901, -2.2623,  ..., -2.3608, -2.3901, -2.3783],\n",
       "        [-2.3099, -2.3961, -2.2686,  ..., -2.3862, -2.3096, -2.4644]],\n",
       "       grad_fn=<SubBackward0>)"
      ]
     },
     "execution_count": 82,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_softmax(pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "caf78754",
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_pred = log_softmax(pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "77044729",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-2.2446, -2.4244, -2.4216,  ..., -2.2878, -2.4068, -2.3208],\n",
       "        [-2.3862, -2.3740, -2.3016,  ..., -2.2498, -2.3517, -2.4064],\n",
       "        [-2.3970, -2.4087, -2.2812,  ..., -2.3484, -2.4407, -2.3838],\n",
       "        ...,\n",
       "        [-2.3303, -2.4678, -2.2690,  ..., -2.4128, -2.2866, -2.4098],\n",
       "        [-2.2764, -2.4901, -2.2623,  ..., -2.3608, -2.3901, -2.3783],\n",
       "        [-2.3099, -2.3961, -2.2686,  ..., -2.3862, -2.3096, -2.4644]],\n",
       "       grad_fn=<SubBackward0>)"
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sm_pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "id": "74866805",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-2.2473, -2.3862, -2.2561], grad_fn=<IndexBackward0>)"
      ]
     },
     "execution_count": 101,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sm_pred[[0, 1, 2], y_train[:3]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "dc4c6e5b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([5, 0, 4])"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_train[:3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "id": "0820acfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def nll(input, target):\n",
    "    return -(input[range(input.shape[0]), target]).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "id": "4283023d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.3178, grad_fn=<NegBackward0>)"
      ]
     },
     "execution_count": 104,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = nll(sm_pred, y_train)\n",
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "id": "ef646892",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.3178, grad_fn=<NllLossBackward0>)"
      ]
     },
     "execution_count": 107,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.nll_loss(F.log_softmax(pred, 1), y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "id": "4a71af4a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.3178, grad_fn=<NllLossBackward0>)"
      ]
     },
     "execution_count": 108,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.cross_entropy(pred, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "id": "c8f8ac67",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_func = F.cross_entropy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "id": "e59af2f7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.0271, -0.1528, -0.1500, -0.0052, -0.0646,  0.0244,  0.1677, -0.0161,\n",
       "        -0.1351, -0.0491], grad_fn=<SelectBackward0>)"
      ]
     },
     "execution_count": 113,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 50\n",
    "xb = x_train[:bs]\n",
    "yb = y_train[:bs]\n",
    "preds = model(xb)\n",
    "preds[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "id": "9dbd19e4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.3221, grad_fn=<NllLossBackward0>)"
      ]
     },
     "execution_count": 114,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss_func(preds, yb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "id": "0cd78b22",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([5, 0, 4, 1, 9, 2, 1, 3, 1, 4, 3, 5, 3, 6, 1, 7, 2, 8, 6, 9, 4, 0, 9, 1,\n",
       "        1, 2, 4, 3, 2, 7, 3, 8, 6, 9, 0, 5, 6, 0, 7, 6, 1, 8, 7, 9, 3, 9, 8, 5,\n",
       "        9, 3])"
      ]
     },
     "execution_count": 115,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "yb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "id": "251becd1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 4, 6, 6, 5, 6, 2, 6, 6, 6, 6, 6,\n",
       "        4, 6, 5, 6, 6, 4, 6, 4, 6, 6, 6, 6, 6, 6, 5, 6, 6, 6, 6, 6, 6, 6, 6, 4,\n",
       "        6, 6])"
      ]
     },
     "execution_count": 116,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "preds.argmax(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "id": "864c2a46",
   "metadata": {},
   "outputs": [],
   "source": [
    "def accuracy(out, yb):\n",
    "    return (out.argmax(1) == yb).float().mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "id": "e6bb1bd5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0600)"
      ]
     },
     "execution_count": 118,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "accuracy(preds, yb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "id": "4ebdbd63",
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 0.5\n",
    "epochs = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "id": "a699089b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def report(loss, preds, yb): \n",
    "    print(f'loss: {loss:.2f}, accuracy: {accuracy(preds, yb):.2f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "id": "4170a6c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss: 2.32, accuracy: 0.06\n"
     ]
    }
   ],
   "source": [
    "xb, yb = x_train[:bs], y_train[:bs]\n",
    "preds = model(xb)\n",
    "report(loss_func(preds, yb), preds, yb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "id": "34376aea",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss: 0.05, accuracy: 0.98\n",
      "loss: 0.05, accuracy: 0.98\n",
      "loss: 0.05, accuracy: 0.98\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(epochs):\n",
    "    for i in range(0, m, bs):\n",
    "        s = slice(i, min(m, i+bs))\n",
    "        xb,yb = x_train[s],y_train[s]\n",
    "        preds = model(xb)\n",
    "        loss = loss_func(preds, yb)\n",
    "        loss.backward()\n",
    "        with torch.no_grad():\n",
    "            for l in model.layers:\n",
    "                if hasattr(l, 'weight'):\n",
    "                    l.weight -= l.weight.grad * lr\n",
    "                    l.bias == l.bias.grad * lr\n",
    "                    l.weight.grad.zero_()\n",
    "                    l.bias.grad.zero_()\n",
    "    report(loss, preds, yb)     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "id": "ed0337d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "m1 = nn.Module()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "f7027598",
   "metadata": {},
   "outputs": [],
   "source": [
    "m1.foo = nn.Linear(3, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "id": "053cf61d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Module(\n",
       "  (foo): Linear(in_features=3, out_features=4, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 137,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "id": "e7e69b9a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('foo', Linear(in_features=3, out_features=4, bias=True))]"
      ]
     },
     "execution_count": 140,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(m1.named_children())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "id": "a51160f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Parameter containing:\n",
       " tensor([[ 0.3285, -0.1131,  0.0071],\n",
       "         [ 0.1433,  0.1641, -0.1538],\n",
       "         [ 0.0191, -0.2399,  0.5618],\n",
       "         [-0.3817,  0.3038, -0.5740]], requires_grad=True),\n",
       " Parameter containing:\n",
       " tensor([-0.0576,  0.4018, -0.1471,  0.0139], requires_grad=True)]"
      ]
     },
     "execution_count": 141,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(m1.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "id": "88321999",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(nn.Module):\n",
    "    def __init__(self, nin, nh, nout):\n",
    "        super().__init__()\n",
    "        self.l1 = nn.Linear(nin, nh)\n",
    "        self.relu = nn.ReLU(nh)\n",
    "        self.l2 = nn.Linear(nh, nout)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.l2(self.relu(self.l1(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "id": "ca23d7b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MLP(n, nh, nout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 162,
   "id": "f56535cf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MLP(\n",
       "  (l1): Linear(in_features=784, out_features=50, bias=True)\n",
       "  (relu): ReLU(inplace=True)\n",
       "  (l2): Linear(in_features=50, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 162,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "id": "223f8bf9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "name: l1, Linear(in_features=784, out_features=50, bias=True)\n",
      "name: relu, ReLU(inplace=True)\n",
      "name: l2, Linear(in_features=50, out_features=10, bias=True)\n"
     ]
    }
   ],
   "source": [
    "for name, l in model.named_children():\n",
    "    print(f'name: {name}, {l}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 164,
   "id": "1aa55500",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([50, 784])\n",
      "torch.Size([50])\n",
      "torch.Size([10, 50])\n",
      "torch.Size([10])\n"
     ]
    }
   ],
   "source": [
    "for p in model.parameters():\n",
    "    print(p.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "id": "f5a13eb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit():\n",
    "    for epoch in range(epochs):\n",
    "        for i in range(0, m, bs):\n",
    "            s = slice(i, min(m, i+bs))\n",
    "            xb,yb = x_train[s], y_train[s]\n",
    "            preds = model(xb)\n",
    "            loss = loss_func(preds, yb)\n",
    "            loss.backward()\n",
    "            with torch.no_grad():\n",
    "                for p in model.parameters():\n",
    "                    p -= p.grad * lr\n",
    "                model.zero_grad()\n",
    "        report(loss, preds, yb)                 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "id": "f09b6a97",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss: 0.08, accuracy: 0.98\n",
      "loss: 0.04, accuracy: 1.00\n",
      "loss: 0.02, accuracy: 1.00\n"
     ]
    }
   ],
   "source": [
    "fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "id": "319ce3e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyModule:\n",
    "    def __init__(self, nin, nh, nout):\n",
    "        self._modules = {}\n",
    "        self.l1 = nn.Linear(nin, nh)\n",
    "        self.relu = nn.ReLU(nh)\n",
    "        self.l2 = nn.Linear(nh, nout)\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",
    "            yield from l.parameters()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 186,
   "id": "ee09b071",
   "metadata": {},
   "outputs": [],
   "source": [
    "mdl = MyModule(n, nh, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "id": "183b9a87",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'l1': Linear(in_features=784, out_features=50, bias=True), 'relu': ReLU(inplace=True), 'l2': Linear(in_features=50, out_features=10, bias=True)}"
      ]
     },
     "execution_count": 187,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mdl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 188,
   "id": "8d9b1f68",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([50, 784])\n",
      "torch.Size([50])\n",
      "torch.Size([10, 50])\n",
      "torch.Size([10])\n"
     ]
    }
   ],
   "source": [
    "for p in mdl.parameters():\n",
    "    print(p.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 191,
   "id": "11399651",
   "metadata": {},
   "outputs": [],
   "source": [
    "layers = [nn.Linear(n, nh), nn.ReLU(nh), nn.Linear(nh, nout)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "id": "76422f4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self, layers):\n",
    "        super().__init__()\n",
    "        self.layers = layers\n",
    "        for i, l in enumerate(self.layers):\n",
    "            self.add_module(f'layer_{i}', l)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        for l in self.layers:\n",
    "            x = l(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 199,
   "id": "d9ad54a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model(layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 200,
   "id": "93194472",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Model(\n",
       "  (layer_0): Linear(in_features=784, out_features=50, bias=True)\n",
       "  (layer_1): ReLU(inplace=True)\n",
       "  (layer_2): Linear(in_features=50, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 200,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 202,
   "id": "1f12d297",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([50, 10])"
      ]
     },
     "execution_count": 202,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(xb).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 203,
   "id": "d7231aff",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import reduce"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 204,
   "id": "07c80639",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self, layers):\n",
    "        super().__init__()\n",
    "        self.layers = layers\n",
    "        for i, l in enumerate(self.layers):\n",
    "            self.add_module(f'layer_{i}', l)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return reduce(lambda x, layer: layer(x), self.layers, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 205,
   "id": "a7af93bf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Model(\n",
       "  (layer_0): Linear(in_features=784, out_features=50, bias=True)\n",
       "  (layer_1): ReLU(inplace=True)\n",
       "  (layer_2): Linear(in_features=50, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 205,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = Model(layers)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 206,
   "id": "50e4d35e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([50, 10])"
      ]
     },
     "execution_count": 206,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(xb).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36c94114",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 222,
   "id": "51670451",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SequentialModel(nn.Module):\n",
    "    def __init__(self, layers):\n",
    "        super().__init__()\n",
    "        self.layers = nn.ModuleList(layers)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        for l in self.layers:\n",
    "            x = l(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 227,
   "id": "1ce92de8",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SequentialModel(layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 228,
   "id": "719af670",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SequentialModel(\n",
       "  (layers): ModuleList(\n",
       "    (0): Linear(in_features=784, out_features=50, bias=True)\n",
       "    (1): ReLU(inplace=True)\n",
       "    (2): Linear(in_features=50, out_features=10, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 228,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 229,
   "id": "2805c185",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([50, 10])"
      ]
     },
     "execution_count": 229,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(xb).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 230,
   "id": "3cf7c6d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss: 0.08, accuracy: 0.96\n",
      "loss: 0.05, accuracy: 0.98\n",
      "loss: 0.05, accuracy: 0.98\n"
     ]
    }
   ],
   "source": [
    "fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 232,
   "id": "aea53d4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = nn.Sequential(*layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 233,
   "id": "d5180523",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): Linear(in_features=784, out_features=50, bias=True)\n",
       "  (1): ReLU(inplace=True)\n",
       "  (2): Linear(in_features=50, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 233,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 234,
   "id": "f097553b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss: 0.15, accuracy: 0.94\n",
      "loss: 0.00, accuracy: 1.00\n",
      "loss: 0.06, accuracy: 0.98\n"
     ]
    }
   ],
   "source": [
    "fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 246,
   "id": "76d13f72",
   "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",
    "        with torch.no_grad():\n",
    "            for p in self.params:\n",
    "                p -= self.lr * p.grad\n",
    "    \n",
    "    def zero_grad(self):\n",
    "        for p in self.params:\n",
    "            p.grad.data.zero_()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 247,
   "id": "c7648606",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = nn.Sequential(*[nn.Linear(n, nh), nn.ReLU(nh), nn.Linear(nh, nout)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 248,
   "id": "ba56e1e2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): Linear(in_features=784, out_features=50, bias=True)\n",
       "  (1): ReLU(inplace=True)\n",
       "  (2): Linear(in_features=50, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 248,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 249,
   "id": "3f1678b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "optim = Optimizer(model.parameters(), lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 250,
   "id": "6ead9df0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss: 0.11, accuracy: 0.97\n",
      "loss: 0.07, accuracy: 0.98\n",
      "loss: 0.06, accuracy: 0.99\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(epochs):\n",
    "    for i in range(0, m, bs):\n",
    "        s = slice(0, min(m, i+bs))\n",
    "        xb, yb = x_train[s], y_train[s]\n",
    "        preds = model(xb)\n",
    "        loss = loss_func(preds, yb)\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "        optim.zero_grad()\n",
    "    report(loss, preds, yb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3086a3b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 262,
   "id": "6f9fae59",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model():\n",
    "    model = nn.Sequential(*[nn.Linear(n, nh), nn.ReLU(nh), nn.Linear(nh, 10)])\n",
    "    optim = Optimizer(model.parameters(), lr)\n",
    "    return model, optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 263,
   "id": "e5aca5d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "model, optim = get_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 264,
   "id": "c41f49fa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.3182, grad_fn=<NllLossBackward0>)"
      ]
     },
     "execution_count": 264,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss_func(model(xb), yb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f52ac90a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 254,
   "id": "e8014adc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dataset():\n",
    "    def __init__(self, x, y):\n",
    "        self.x, self.y = x, 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": 255,
   "id": "2cc3d2ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds, val_ds = Dataset(x_train, y_train), Dataset(x_val, y_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 256,
   "id": "fd336313",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(train_ds) == len(x_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 257,
   "id": "f3d0059f",
   "metadata": {},
   "outputs": [],
   "source": [
    "xb, yb = train_ds[:50]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 259,
   "id": "9a3134b6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([50, 784]), torch.Size([50]))"
      ]
     },
     "execution_count": 259,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xb.shape, yb.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 265,
   "id": "67e7be81",
   "metadata": {},
   "outputs": [],
   "source": [
    "model, optim = get_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 266,
   "id": "8ecc5fbd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss: 0.12, accuracy: 0.96\n",
      "loss: 0.08, accuracy: 0.98\n",
      "loss: 0.05, accuracy: 1.00\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(epochs):\n",
    "    for i in range(0, m, bs):\n",
    "        xb, yb = train_ds[i:i+bs]\n",
    "        preds = model(xb)\n",
    "        loss = loss_func(preds, yb)\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "        optim.zero_grad()\n",
    "    report(loss, preds, yb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36c69962",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 273,
   "id": "43216325",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataLoader():\n",
    "    def __init__(self, ds, bs):\n",
    "        self.ds, self.bs = ds, bs\n",
    "    \n",
    "    def __iter__(self):\n",
    "        for i in range(0, len(self.ds), self.bs):\n",
    "            yield self.ds[i:i+self.bs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 274,
   "id": "919cb4ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dl = DataLoader(train_ds, bs)\n",
    "val_dl = DataLoader(val_ds, bs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 275,
   "id": "322b71c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "xb, yb = next(iter(train_dl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 276,
   "id": "0ca8f7bf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([50, 784]), torch.Size([50]))"
      ]
     },
     "execution_count": 276,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xb.shape, yb.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 278,
   "id": "392f9b2a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(5)"
      ]
     },
     "execution_count": 278,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD4CAYAAAAq5pAIAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAAsTAAALEwEAmpwYAAAN80lEQVR4nO3df6hcdXrH8c+ncf3DrBpTMYasNhuRWBWbLRqLSl2RrD9QNOqWDVgsBrN/GHChhEr6xyolEuqP0qAsuYu6sWyzLqgYZVkVo6ZFCF5j1JjU1YrdjV6SSozG+KtJnv5xT+Su3vnOzcyZOZP7vF9wmZnzzJnzcLife87Md879OiIEYPL7k6YbANAfhB1IgrADSRB2IAnCDiRxRD83ZpuP/oEeiwiPt7yrI7vtS22/aftt27d281oAesudjrPbniLpd5IWSNou6SVJiyJia2EdjuxAj/XiyD5f0tsR8U5EfCnpV5Ku6uL1APRQN2GfJekPYx5vr5b9EdtLbA/bHu5iWwC61M0HdOOdKnzjND0ihiQNSZzGA03q5si+XdJJYx5/R9L73bUDoFe6CftLkk61/V3bR0r6kaR19bQFoG4dn8ZHxD7bSyU9JWmKpAci4o3aOgNQq46H3jraGO/ZgZ7ryZdqABw+CDuQBGEHkiDsQBKEHUiCsANJEHYgCcIOJEHYgSQIO5AEYQeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgrADSRB2IAnCDiRB2IEkCDuQBGEHkiDsQBKEHUii4ymbcXiYMmVKsX7sscf2dPtLly5tWTvqqKOK686dO7dYv/nmm4v1u+66q2Vt0aJFxXU///zzYn3lypXF+u23316sN6GrsNt+V9IeSfsl7YuIs+toCkD96jiyXxQRH9TwOgB6iPfsQBLdhj0kPW37ZdtLxnuC7SW2h20Pd7ktAF3o9jT+/Ih43/YJkp6x/V8RsWHsEyJiSNKQJNmOLrcHoENdHdkj4v3qdqekxyTNr6MpAPXrOOy2p9o++uB9ST+QtKWuxgDUq5vT+BmSHrN98HX+PSJ+W0tXk8zJJ59crB955JHF+nnnnVesX3DBBS1r06ZNK6577bXXFutN2r59e7G+atWqYn3hwoUta3v27Cmu++qrrxbrL7zwQrE+iDoOe0S8I+kvauwFQA8x9AYkQdiBJAg7kARhB5Ig7EASjujfl9om6zfo5s2bV6yvX7++WO/1ZaaD6sCBA8X6jTfeWKx/8sknHW97ZGSkWP/www+L9TfffLPjbfdaRHi85RzZgSQIO5AEYQeSIOxAEoQdSIKwA0kQdiAJxtlrMH369GJ948aNxfqcOXPqbKdW7XrfvXt3sX7RRRe1rH355ZfFdbN+/6BbjLMDyRF2IAnCDiRB2IEkCDuQBGEHkiDsQBJM2VyDXbt2FevLli0r1q+44opi/ZVXXinW2/1L5ZLNmzcX6wsWLCjW9+7dW6yfccYZLWu33HJLcV3UiyM7kARhB5Ig7EAShB1IgrADSRB2IAnCDiTB9ewD4JhjjinW200vvHr16pa1xYsXF9e9/vrri/W1a9cW6xg8HV/PbvsB2zttbxmzbLrtZ2y/Vd0eV2ezAOo3kdP4X0i69GvLbpX0bEScKunZ6jGAAdY27BGxQdLXvw96laQ11f01kq6uty0Adev0u/EzImJEkiJixPYJrZ5oe4mkJR1uB0BNen4hTEQMSRqS+IAOaFKnQ287bM+UpOp2Z30tAeiFTsO+TtIN1f0bJD1eTzsAeqXtabzttZK+L+l429sl/VTSSkm/tr1Y0u8l/bCXTU52H3/8cVfrf/TRRx2ve9NNNxXrDz/8cLHebo51DI62YY+IRS1KF9fcC4Ae4uuyQBKEHUiCsANJEHYgCcIOJMElrpPA1KlTW9aeeOKJ4roXXnhhsX7ZZZcV608//XSxjv5jymYgOcIOJEHYgSQIO5AEYQeSIOxAEoQdSIJx9knulFNOKdY3bdpUrO/evbtYf+6554r14eHhlrX77ruvuG4/fzcnE8bZgeQIO5AEYQeSIOxAEoQdSIKwA0kQdiAJxtmTW7hwYbH+4IMPFutHH310x9tevnx5sf7QQw8V6yMjIx1vezJjnB1IjrADSRB2IAnCDiRB2IEkCDuQBGEHkmCcHUVnnnlmsX7PPfcU6xdf3Plkv6tXry7WV6xYUay/9957HW/7cNbxOLvtB2zvtL1lzLLbbL9ne3P1c3mdzQKo30RO438h6dJxlv9LRMyrfn5Tb1sA6tY27BGxQdKuPvQCoIe6+YBuqe3XqtP841o9yfYS28O2W/8zMgA912nYfybpFEnzJI1IurvVEyNiKCLOjoizO9wWgBp0FPaI2BER+yPigKSfS5pfb1sA6tZR2G3PHPNwoaQtrZ4LYDC0HWe3vVbS9yUdL2mHpJ9Wj+dJCknvSvpxRLS9uJhx9sln2rRpxfqVV17ZstbuWnl73OHir6xfv75YX7BgQbE+WbUaZz9iAisuGmfx/V13BKCv+LoskARhB5Ig7EAShB1IgrADSXCJKxrzxRdfFOtHHFEeLNq3b1+xfskll7SsPf/888V1D2f8K2kgOcIOJEHYgSQIO5AEYQeSIOxAEoQdSKLtVW/I7ayzzirWr7vuumL9nHPOaVlrN47eztatW4v1DRs2dPX6kw1HdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgnH2SW7u3LnF+tKlS4v1a665plg/8cQTD7mnidq/f3+xPjJS/u/lBw4cqLOdwx5HdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgnH2w0C7sexFi8abaHdUu3H02bNnd9JSLYaHh4v1FStWFOvr1q2rs51Jr+2R3fZJtp+zvc32G7ZvqZZPt/2M7beq2+N63y6ATk3kNH6fpL+PiD+X9FeSbrZ9uqRbJT0bEadKerZ6DGBAtQ17RIxExKbq/h5J2yTNknSVpDXV09ZIurpHPQKowSG9Z7c9W9L3JG2UNCMiRqTRPwi2T2ixzhJJS7rsE0CXJhx229+W9Iikn0TEx/a4c8d9Q0QMSRqqXoOJHYGGTGjozfa3NBr0X0bEo9XiHbZnVvWZknb2pkUAdWh7ZPfoIfx+Sdsi4p4xpXWSbpC0srp9vCcdTgIzZswo1k8//fRi/d577y3WTzvttEPuqS4bN24s1u+8886WtccfL//KcIlqvSZyGn++pL+V9LrtzdWy5RoN+a9tL5b0e0k/7EmHAGrRNuwR8Z+SWr1Bv7jedgD0Cl+XBZIg7EAShB1IgrADSRB2IAkucZ2g6dOnt6ytXr26uO68efOK9Tlz5nTSUi1efPHFYv3uu+8u1p966qli/bPPPjvkntAbHNmBJAg7kARhB5Ig7EAShB1IgrADSRB2IIk04+znnntusb5s2bJiff78+S1rs2bN6qinunz66acta6tWrSque8cddxTre/fu7agnDB6O7EAShB1IgrADSRB2IAnCDiRB2IEkCDuQRJpx9oULF3ZV78bWrVuL9SeffLJY37dvX7FeuuZ89+7dxXWRB0d2IAnCDiRB2IEkCDuQBGEHkiDsQBKEHUjCEVF+gn2SpIcknSjpgKShiPhX27dJuknS/1ZPXR4Rv2nzWuWNAehaRIw76/JEwj5T0syI2GT7aEkvS7pa0t9I+iQi7ppoE4Qd6L1WYZ/I/Owjkkaq+3tsb5PU7L9mAXDIDuk9u+3Zkr4naWO1aKnt12w/YPu4FusssT1se7i7VgF0o+1p/FdPtL8t6QVJKyLiUdszJH0gKST9k0ZP9W9s8xqcxgM91vF7dkmy/S1JT0p6KiLuGac+W9KTEXFmm9ch7ECPtQp729N425Z0v6RtY4NefXB30EJJW7ptEkDvTOTT+Ask/Yek1zU69CZJyyUtkjRPo6fx70r6cfVhXum1OLIDPdbVaXxdCDvQex2fxgOYHAg7kARhB5Ig7EAShB1IgrADSRB2IAnCDiRB2IEkCDuQBGEHkiDsQBKEHUiCsANJ9HvK5g8k/c+Yx8dXywbRoPY2qH1J9NapOnv7s1aFvl7P/o2N28MRcXZjDRQMam+D2pdEb53qV2+cxgNJEHYgiabDPtTw9ksGtbdB7Uuit071pbdG37MD6J+mj+wA+oSwA0k0Enbbl9p+0/bbtm9toodWbL9r+3Xbm5uen66aQ2+n7S1jlk23/Yztt6rbcefYa6i322y/V+27zbYvb6i3k2w/Z3ub7Tds31Itb3TfFfrqy37r+3t221Mk/U7SAknbJb0kaVFEbO1rIy3YflfS2RHR+BcwbP+1pE8kPXRwai3b/yxpV0SsrP5QHhcR/zAgvd2mQ5zGu0e9tZpm/O/U4L6rc/rzTjRxZJ8v6e2IeCcivpT0K0lXNdDHwIuIDZJ2fW3xVZLWVPfXaPSXpe9a9DYQImIkIjZV9/dIOjjNeKP7rtBXXzQR9lmS/jDm8XYN1nzvIelp2y/bXtJ0M+OYcXCarer2hIb7+bq203j309emGR+YfdfJ9OfdaiLs401NM0jjf+dHxF9KukzSzdXpKibmZ5JO0egcgCOS7m6ymWqa8Uck/SQiPm6yl7HG6asv+62JsG+XdNKYx9+R9H4DfYwrIt6vbndKekyjbzsGyY6DM+hWtzsb7ucrEbEjIvZHxAFJP1eD+66aZvwRSb+MiEerxY3vu/H66td+ayLsL0k61fZ3bR8p6UeS1jXQxzfYnlp9cCLbUyX9QIM3FfU6STdU92+Q9HiDvfyRQZnGu9U042p43zU+/XlE9P1H0uUa/UT+vyX9YxM9tOhrjqRXq583mu5N0lqNntb9n0bPiBZL+lNJz0p6q7qdPkC9/ZtGp/Z+TaPBmtlQbxdo9K3ha5I2Vz+XN73vCn31Zb/xdVkgCb5BByRB2IEkCDuQBGEHkiDsQBKEHUiCsANJ/D+f1mbt6t55/AAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(xb[0].view(28, 28))\n",
    "yb[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 282,
   "id": "c3d0d8e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "model, optim = get_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 283,
   "id": "afc0c960",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit():\n",
    "    for epoch in range(epochs):\n",
    "        for xb, yb, in train_dl:\n",
    "            preds = model(xb)\n",
    "            loss = loss_func(preds, yb)\n",
    "            loss.backward()\n",
    "            \n",
    "            optim.step()\n",
    "            optim.zero_grad()\n",
    "        report(loss, preds, yb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 284,
   "id": "e82a0660",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss: 0.15, accuracy: 0.96\n",
      "loss: 0.08, accuracy: 0.96\n",
      "loss: 0.07, accuracy: 0.96\n"
     ]
    }
   ],
   "source": [
    "fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 290,
   "id": "c97501f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from itertools import islice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 318,
   "id": "5ec92565",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sampler:\n",
    "    def __init__(self, ds, shuffle=False):\n",
    "        self.n, self.shuffle = len(ds), shuffle\n",
    "    \n",
    "    def __iter__(self):\n",
    "        res = list(range(self.n))\n",
    "        if self.shuffle:\n",
    "            random.shuffle(res)\n",
    "        return iter(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 319,
   "id": "91a3aeef",
   "metadata": {},
   "outputs": [],
   "source": [
    "ss = Sampler(train_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 320,
   "id": "544073bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "it = iter(ss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 321,
   "id": "1fe93469",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "1\n",
      "2\n",
      "3\n",
      "4\n"
     ]
    }
   ],
   "source": [
    "for ele in range(5):\n",
    "    print(next(it))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 322,
   "id": "7c2bcf65",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0, 1, 2, 3, 4]"
      ]
     },
     "execution_count": 322,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(islice(ss, 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 323,
   "id": "04bca889",
   "metadata": {},
   "outputs": [],
   "source": [
    "import fastcore.all as fc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 324,
   "id": "ed3ebeff",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BatchSampler():\n",
    "    def __init__(self, sampler, bs, drop_last=False):\n",
    "        fc.store_attr()\n",
    "    \n",
    "    def __iter__(self):\n",
    "        yield from fc.chunked(iter(self.sampler), self.bs, drop_last= self.drop_last)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 325,
   "id": "c748bf7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "batches = BatchSampler(ss, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 326,
   "id": "1066ecce",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[0, 1, 2, 3],\n",
       " [4, 5, 6, 7],\n",
       " [8, 9, 10, 11],\n",
       " [12, 13, 14, 15],\n",
       " [16, 17, 18, 19]]"
      ]
     },
     "execution_count": 326,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(islice(batches, 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 327,
   "id": "86368ef6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate(b):\n",
    "    xs, ys = zip(*b)\n",
    "    return torch.stack(xs), torch.stack(ys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 332,
   "id": "371951e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataLoader():\n",
    "    def __init__(self, ds, batches, collate_fn=collate):\n",
    "        fc.store_attr()\n",
    "    \n",
    "    def __iter__(self):\n",
    "        yield from (self.collate_fn(self.ds[i] for i in b) for b in self.batches)       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 333,
   "id": "4196be41",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_samp = BatchSampler(Sampler(train_ds, shuffle=True), bs)\n",
    "val_samp = BatchSampler(Sampler(val_ds, shuffle=True), bs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 334,
   "id": "4df9fbbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dl = DataLoader(train_ds, train_samp, collate_fn=collate)\n",
    "val_dl = DataLoader(val_ds, val_samp, collate_fn=collate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 343,
   "id": "fc75ac85",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([784])\n"
     ]
    }
   ],
   "source": [
    "print(next(iter(train_dl))[0][0].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 344,
   "id": "6d526283",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x13ea56110>"
      ]
     },
     "execution_count": 344,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD4CAYAAAAq5pAIAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAAsTAAALEwEAmpwYAAANs0lEQVR4nO3db6xU9Z3H8c9HaUMCPABRA7e4QOPfmAgbYjSSTdem6KoJ9kE3JWpYagKakpRkHyypmppsmpjNwj6zEaIpu6INRioENK0hzeL6oPEqVKBIZQ1b6CWAwVBBI6t898E9dC9wz5nLzJk5A9/3K7mZmfOdM+ebCR9+Z+acOT9HhABc/q5ougEAvUHYgSQIO5AEYQeSIOxAEuN6uTHbfPUPdFlEeLTlHY3stu+1vc/2ftsrO3ktAN3ldo+z275S0h8kfUfSIUnvSFoUEb+vWIeRHeiybozst0vaHxEfRcRpSb+QtLCD1wPQRZ2EfUDSwRGPDxXLzmF7qe1B24MdbAtAhzr5gm60XYULdtMjYo2kNRK78UCTOhnZD0maMeLxNyQNddYOgG7pJOzvSLre9izbX5f0fUmb62kLQN3a3o2PiC9tL5f0K0lXSnohIvbU1hmAWrV96K2tjfGZHei6rpxUA+DSQdiBJAg7kARhB5Ig7EAShB1IgrADSRB2IAnCDiRB2IEkCDuQBGEHkiDsQBKEHUiCsANJEHYgCcIOJEHYgSQIO5AEYQeSIOxAEj2dshnd8fDDD5fW1q1b19Frr1q1qrLe6vX37OHq4v2CkR1IgrADSRB2IAnCDiRB2IEkCDuQBGEHkmAW10vAkiVLKuurV68urU2aNKnuds4xNDRUWf/kk09Ka/v3769c97HHHqusHzt2rLKeVdksrh2dVGP7gKRPJX0l6cuImNfJ6wHonjrOoPvbiPi4htcB0EV8ZgeS6DTsIenXtt+1vXS0J9heanvQ9mCH2wLQgU534++KiCHb10h60/YHEbF95BMiYo2kNRJf0AFN6mhkj4ih4vaopF9Kur2OpgDUr+2w255ge9LZ+5IWSNpdV2MA6tX2cXbbszU8mkvDHwdeioiftliH3fhRjB8/vrK+b9++yvrAwECd7dTKHvWQrySp1b+9bdu2Vdbvueeetnq63NV+nD0iPpJ0W9sdAegpDr0BSRB2IAnCDiRB2IEkCDuQBJeS7oFWh9Y2btxYWe/nQ2vdNG8eP6KsEyM7kARhB5Ig7EAShB1IgrADSRB2IAnCDiTBcfYaTJkypbK+du3ayvqCBQvqbOeyMW5c9T/Pm266qbL+wQcf1NnOJY+RHUiCsANJEHYgCcIOJEHYgSQIO5AEYQeS4Dj7GN1///2ltZUrV1aue+edd9bdTm1OnTpVWX/iiScq663OEah631qZMGFCZf2hhx6qrD/11FNtb/tyxMgOJEHYgSQIO5AEYQeSIOxAEoQdSIKwA0lwnH2MrrvuutJaPx9Hf+ONNyrrzz33XGV9y5YtlfWbb775ontCM1qO7LZfsH3U9u4Ry6bYftP2h8Xt5O62CaBTY9mN/7mke89btlLStoi4XtK24jGAPtYy7BGxXdLx8xYvlLSuuL9O0oP1tgWgbu1+Zr82Ig5LUkQctn1N2RNtL5W0tM3tAKhJ17+gi4g1ktZIku3o9vYAjK7dQ29HbE+TpOL2aH0tAeiGdsO+WdLi4v5iSZvqaQdAt7Tcjbf9sqRvSZpq+5Ckn0h6RtIG249K+qOk73WzSVSrmt/9kUceqVz3iy++qLud2pw+fbqyvmvXrh51cnloGfaIWFRS+nbNvQDoIk6XBZIg7EAShB1IgrADSRB2IAl+4jpGr732WmntgQceqFx3+vTplfWpU6dW1t9+++3K+uOPP15a6+dDa618/vnnlfUNGzb0qJPLAyM7kARhB5Ig7EAShB1IgrADSRB2IAnCDiThiN5dPIYr1Yxuzpw5lfX9+/dX1k+ePFljN+caGBiorL/++uuV9VtvvbW01urf3okTJyrrV111VWU9q4jwaMsZ2YEkCDuQBGEHkiDsQBKEHUiCsANJEHYgCX7P3gd27tzZ2LZvuOGGyvrWrVsr67Nnz66sX3FF+Xhy5syZynXfeuutyjouDiM7kARhB5Ig7EAShB1IgrADSRB2IAnCDiTBcfbL3MyZMyvrmzdvrqzPmjWrst7qN+lVx9JbXfe96nr4uHgtR3bbL9g+anv3iGVP2/6T7Z3F333dbRNAp8ayG/9zSfeOsvzfImJO8Vd9uRIAjWsZ9ojYLul4D3oB0EWdfEG33Pb7xW7+5LIn2V5qe9D2YAfbAtChdsP+M0nflDRH0mFJq8qeGBFrImJeRMxrc1sAatBW2CPiSER8FRFnJK2VdHu9bQGoW1thtz1txMPvStpd9lwA/aHldeNtvyzpW5KmSjoi6SfF4zmSQtIBScsi4nDLjXHd+K6YMWNGaa3V79FvueWWuts5hz3qJcwlSbfddlvlurt3M4a0o+y68S1PqomIRaMsfr7jjgD0FKfLAkkQdiAJwg4kQdiBJAg7kAQ/cb0EtPqp5/Lly0trN954Y93tXJQdO3aU1o4dO1a57vjx4yvrK1asaKelnnjxxRcr64cOHepRJ/+PkR1IgrADSRB2IAnCDiRB2IEkCDuQBGEHkmj5E9daN5b0J65z5syprC9ZsqSyvmzZssr6uHH9e7rEqVOnSmvHj1df2rBqumdJGhgYaKunXti+fXtl/e677+7atst+4srIDiRB2IEkCDuQBGEHkiDsQBKEHUiCsANJ9O8B2kvI/PnzK+svvfRSZX369Ol1ttNXJk6cWFqbMGFCDzvprcmTS2dEawwjO5AEYQeSIOxAEoQdSIKwA0kQdiAJwg4kwXH2MbrjjjtKa6+88krluldffXXd7aDPbdq0qekWLtByZLc9w/ZvbO+1vcf2j4rlU2y/afvD4rb/ziIA8Bdj2Y3/UtI/RsTNku6Q9EPbt0haKWlbRFwvaVvxGECfahn2iDgcEe8V9z+VtFfSgKSFktYVT1sn6cEu9QigBhf1md32TElzJf1W0rURcVga/g/B9jUl6yyVtLTDPgF0aMxhtz1R0quSVkTEn+1Rr2l3gYhYI2lN8RopLzgJ9IMxHXqz/TUNB319RGwsFh+xPa2oT5N0tDstAqhDy5Hdw0P485L2RsTqEaXNkhZLeqa47b9jDTWaO3duaY1Da+3ZunVrZf3EiRM96uTiDQ4OVtafffbZHnUydmPZjb9L0iOSdtneWSz7sYZDvsH2o5L+KOl7XekQQC1ahj0i/ktS2Qf0b9fbDoBu4XRZIAnCDiRB2IEkCDuQBGEHkuAnrpeB06dPl9b27dvXw04utHz58tLajh07Ktf97LPP6m4nNUZ2IAnCDiRB2IEkCDuQBGEHkiDsQBKEHUiC4+xjtH379tLa0NBQ5bqdTsn85JNPVtYPHjxYWlu/fn1H28blg5EdSIKwA0kQdiAJwg4kQdiBJAg7kARhB5JwRO8maWFGGKD7ImLUq0EzsgNJEHYgCcIOJEHYgSQIO5AEYQeSIOxAEi3DbnuG7d/Y3mt7j+0fFcuftv0n2zuLv/u63y6AdrU8qcb2NEnTIuI925MkvSvpQUl/L+lkRPzrmDfGSTVA15WdVDOW+dkPSzpc3P/U9l5JA/W2B6DbLuozu+2ZkuZK+m2xaLnt922/YHtyyTpLbQ/aHuysVQCdGPO58bYnSvpPST+NiI22r5X0saSQ9M8a3tX/QYvXYDce6LKy3fgxhd321yRtkfSriFg9Sn2mpC0RcWuL1yHsQJe1/UMY25b0vKS9I4NefHF31ncl7e60SQDdM5Zv4+dLekvSLklnisU/lrRI0hwN78YfkLSs+DKv6rUY2YEu62g3vi6EHeg+fs8OJEfYgSQIO5AEYQeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IouUFJ2v2saT/GfF4arGsH/Vrb/3al0Rv7aqzt78qK/T09+wXbNwejIh5jTVQoV9769e+JHprV696YzceSIKwA0k0HfY1DW+/Sr/21q99SfTWrp701uhndgC90/TIDqBHCDuQRCNht32v7X2299te2UQPZWwfsL2rmIa60fnpijn0jtrePWLZFNtv2v6wuB11jr2GeuuLabwrphlv9L1revrznn9mt32lpD9I+o6kQ5LekbQoIn7f00ZK2D4gaV5ENH4Chu2/kXRS0r+fnVrL9r9IOh4RzxT/UU6OiH/qk96e1kVO492l3sqmGf8HNfje1Tn9eTuaGNlvl7Q/Ij6KiNOSfiFpYQN99L2I2C7p+HmLF0paV9xfp+F/LD1X0ltfiIjDEfFecf9TSWenGW/0vavoqyeaCPuApIMjHh9Sf833HpJ+bftd20ubbmYU156dZqu4vabhfs7XchrvXjpvmvG+ee/amf68U02EfbSpafrp+N9dEfHXkv5O0g+L3VWMzc8kfVPDcwAelrSqyWaKacZflbQiIv7cZC8jjdJXT963JsJ+SNKMEY+/IWmogT5GFRFDxe1RSb/U8MeOfnLk7Ay6xe3Rhvv5i4g4EhFfRcQZSWvV4HtXTDP+qqT1EbGxWNz4ezdaX71635oI+zuSrrc9y/bXJX1f0uYG+riA7QnFFyeyPUHSAvXfVNSbJS0u7i+WtKnBXs7RL9N4l00zrobfu8anP4+Inv9Juk/D38j/t6QnmuihpK/Zkn5X/O1pujdJL2t4t+5/NbxH9KikqyRtk/RhcTulj3r7Dw1P7f2+hoM1raHe5mv4o+H7knYWf/c1/d5V9NWT943TZYEkOIMOSIKwA0kQdiAJwg4kQdiBJAg7kARhB5L4P88DKOvbolrhAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(next(iter(train_dl))[0][0].view(28, 28))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 345,
   "id": "e64e049a",
   "metadata": {},
   "outputs": [],
   "source": [
    "xb, yb = next(iter(train_dl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 346,
   "id": "103d2d91",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([50, 784]), torch.Size([50]))"
      ]
     },
     "execution_count": 346,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xb.shape, yb.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 348,
   "id": "5c0ec660",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0)"
      ]
     },
     "execution_count": 348,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD4CAYAAAAq5pAIAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAAsTAAALEwEAmpwYAAAOOElEQVR4nO3df6hc9ZnH8c/H+AtswWjIJWujMSV/KMZNJQShsihqdf3DRKFSBYm0cBvRUGFhjV2wF2RF180u+AMhRWlca7So3YYqtiaIPwhUrxJjbIxmQ9akXhI0xF9/pBt99o97Uq7Jne/czMyZM8nzfsFl5p5nzjkPh/u558ycM+friBCAY99xTTcAoD8IO5AEYQeSIOxAEoQdSOL4fq7MNh/9AzWLCE82vas9u+0rbW+1vc32im6WBaBe7vQ8u+1pkt6XdLmkXZLekHR9RPy5MA97dqBmdezZF0naFhHbI+Kvkp6UtLiL5QGoUTdhP0PSzgm/76qmfYPtYdujtke7WBeALnXzAd1khwqHHaZHxCpJqyQO44EmdbNn3yVp9oTfvyPpo+7aAVCXbsL+hqR5ts+2faKkH0la25u2APRax4fxEXHA9q2S/iBpmqRHI+LdnnUGoKc6PvXW0cp4zw7UrpaLagAcPQg7kARhB5Ig7EAShB1IgrADSRB2IAnCDiRB2IEkCDuQBGEHkiDsQBKEHUiCsANJEHYgCcIOJEHYgSQIO5AEYQeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgrADSXQ8Prsk2d4h6XNJX0k6EBELe9EUgN7rKuyVSyLi4x4sB0CNOIwHkug27CHpj7bftD082QtsD9setT3a5boAdMER0fnM9t9FxEe2Z0p6UdLyiHil8PrOVwZgSiLCk03vas8eER9Vj3sk/VbSom6WB6A+HYfd9im2v33wuaQfSNrcq8YA9FY3n8YPSfqt7YPLeSIiXuhJVzhqnH766cX61Vdf3adODjd//vyWtRtvvLHWdY+MjBTrDz30UK3rn0zHYY+I7ZL+voe9AKgRp96AJAg7kARhB5Ig7EAShB1Ioqsr6I54ZVxBN3Auu+yyYv3kk08u1pcvX16sX3755Ufc07Hg5ZdfLtYvueSS2tZdyxV0AI4ehB1IgrADSRB2IAnCDiRB2IEkCDuQRC9uOIkBtmhR+X4iTzzxRLE+Y8aMXrbTVxs3bmxZ279/f1fLnjt3brF+5plndrX8OrBnB5Ig7EAShB1IgrADSRB2IAnCDiRB2IEkOM9+FFi2bFmxXvrO+OLFi4vzHndcvf/vDxw40LK2ffv24rx33HFHsd7uXPcjjzzSsrZv377ivENDQ8X6DTfcUKzPmTOnWG8Ce3YgCcIOJEHYgSQIO5AEYQeSIOxAEoQdSIL7xvfBtGnTivWbb765q/o555xzxD1N1ZYtW4r1+++/v1j/4osvWtaeeuqp4rwrV64s1ttdf/D000+3rN11113Fed97771ifZB1fN9424/a3mN784Rpp9l+0fYH1eP0XjYLoPemchj/K0lXHjJthaT1ETFP0vrqdwADrG3YI+IVSXsPmbxY0urq+WpJS3rbFoBe6/Ta+KGIGJOkiBizPbPVC20PSxrucD0AeqT2L8JExCpJq6S8H9ABg6DTU2+7bc+SpOpxT+9aAlCHTsO+VtLS6vlSSb/rTTsA6tL2PLvtNZIuljRD0m5Jv5D035J+I+lMSR9K+mFEHPoh3mTLSnkYf+GFFxbrGzZs6FMnh9u0aVOxfu+99xbra9as6WU733DRRRcV6+vWrSvWTzzxxJa1Dz/8sDjvFVdcUaxv3bq1WG9Sq/Psbd+zR8T1LUqXdtURgL7iclkgCcIOJEHYgSQIO5AEYQeS4CuuPXDuuecW688991yxftZZZ/WynW948MEHi/WRkZFife/etmdUG/P8888X61deeej3t6au3SnHdre5blLHX3EFcGwg7EAShB1IgrADSRB2IAnCDiRB2IEkGLJ5ihYuXNiy9thjjxXnnTVrVq/b+Yb77ruvZW3FivK9QPt5nUWvzZ8/v7Zlt/vq79GIPTuQBGEHkiDsQBKEHUiCsANJEHYgCcIOJMF59ikaHm49gtXZZ59dnLd0S2NJ+uSTT4r1m266qVgfHR1tWRvk8+jHH1/+83vyySeL9W6uX2g3ZHO7exAcjdizA0kQdiAJwg4kQdiBJAg7kARhB5Ig7EASnGevDA0NFesLFixoWTvppJO6WveOHTuK9bfffrtY3717d1frr1PpGoPzzz+/OG9pm0vSccd1vq96//33i/XPPvus42UPqrZby/ajtvfY3jxh2ojtv9jeWP1cVW+bALo1lX+Nv5I02dAa/xkRC6qf8tAcABrXNuwR8YqkwR0DCMCUdPMB3a22N1WH+dNbvcj2sO1R260v4AZQu07D/rCk70paIGlM0spWL4yIVRGxMCJa37ERQO06CntE7I6IryLia0m/lLSot20B6LWOwm574ncLr5G0udVrAQyGtufZba+RdLGkGbZ3SfqFpIttL5AUknZI+ml9LfbH0qVLi/XSfePb2bZtW7F+7bXXFuu7du3qeN1NmzdvXsva66+/3tWyv/zyy2K9dM/8F154oat1H43ahj0irp9k8iM19AKgRlwuCyRB2IEkCDuQBGEHkiDsQBJ8xbUP2t0qeufOnX3qpPduueWWYn358uUdL3v9+vXF+t13312sv/TSSx2v+1jEnh1IgrADSRB2IAnCDiRB2IEkCDuQBGEHkuA8e+XSSy+tbdnr1q2rbdndOu+884r1a665plhv99XguXPnHnFPB7Ub6nrDhg0dLzsj9uxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kATn2SsHDhyobdknnHBCbcuWpEWLWo/RccEFFxTnHRkZKdZnzpzZSUt/U7rd8+OPP16c9/nny+OF7t+/v6OesmLPDiRB2IEkCDuQBGEHkiDsQBKEHUiCsANJOCL6tzK7fys7Qu22Qzfb6c477yzW9+zZU6zffvvtxfr06dNb1k499dTivN169dVXi/WVK1e2rK1du7bX7UBSRHiy6W337LZn237J9hbb79r+WTX9NNsv2v6gemz9FwegcVM5jD8g6Z8i4hxJF0q6xfa5klZIWh8R8yStr34HMKDahj0ixiLirer555K2SDpD0mJJq6uXrZa0pKYeAfTAEV0bb3uOpO9J+pOkoYgYk8b/Idie9CJq28OShrvsE0CXphx229+S9Iyk2yLiM3vSzwAOExGrJK2qljGwH9ABx7opnXqzfYLGg/7riHi2mrzb9qyqPktS+SNlAI1qe+rN47vw1ZL2RsRtE6bfJ+mTiLjH9gpJp0XEP7dZ1sDu2R944IFivd3QxEerdqf9li1bVqy3Gxb5008/PeKe0J1Wp96mchj/fUk3SnrH9sZq2s8l3SPpN7Z/IulDST/sQZ8AatI27BHxmqRWb9DrG1kBQE9xuSyQBGEHkiDsQBKEHUiCsANJcCvpys6dO5tuoaWxsbFifd++fS1rr732WnHehx9+uFjfuHFjsY6jB3t2IAnCDiRB2IEkCDuQBGEHkiDsQBKEHUiC8+yVdsMDz549u2Xtuuuu62rde/fuLdaXLFlSrG/durWr9SMH9uxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kARDNgPHmI6HbAZwbCDsQBKEHUiCsANJEHYgCcIOJEHYgSTaht32bNsv2d5i+13bP6umj9j+i+2N1c9V9bcLoFNtL6qxPUvSrIh4y/a3Jb0paYmk6yR9ERH/PuWVcVENULtWF9VMZXz2MUlj1fPPbW+RdEZv2wNQtyN6z257jqTvSfpTNelW25tsP2p7eot5hm2P2h7trlUA3ZjytfG2vyXpZUn/GhHP2h6S9LGkkHSXxg/1f9xmGRzGAzVrdRg/pbDbPkHS7yX9ISL+Y5L6HEm/j4jz2iyHsAM16/iLMLYt6RFJWyYGvfrg7qBrJG3utkkA9ZnKp/EXSXpV0juSvq4m/1zS9ZIWaPwwfoekn1Yf5pWWxZ4dqFlXh/G9QtiB+vF9diA5wg4kQdiBJAg7kARhB5Ig7EAShB1IgrADSRB2IAnCDiRB2IEkCDuQBGEHkiDsQBJtbzjZYx9L+t8Jv8+opg2iQe1tUPuS6K1TveztrFaFvn6f/bCV26MRsbCxBgoGtbdB7Uuit071qzcO44EkCDuQRNNhX9Xw+ksGtbdB7Uuit071pbdG37MD6J+m9+wA+oSwA0k0EnbbV9reanub7RVN9NCK7R2236mGoW50fLpqDL09tjdPmHaa7Rdtf1A9TjrGXkO9DcQw3oVhxhvddk0Pf9739+y2p0l6X9LlknZJekPS9RHx57420oLtHZIWRkTjF2DY/gdJX0h67ODQWrb/TdLeiLin+kc5PSJuH5DeRnSEw3jX1FurYcZvUoPbrpfDn3eiiT37IknbImJ7RPxV0pOSFjfQx8CLiFck7T1k8mJJq6vnqzX+x9J3LXobCBExFhFvVc8/l3RwmPFGt12hr75oIuxnSNo54fddGqzx3kPSH22/aXu46WYmMXRwmK3qcWbD/Ryq7TDe/XTIMOMDs+06Gf68W02EfbKhaQbp/N/3I+ICSf8o6ZbqcBVT87Ck72p8DMAxSSubbKYaZvwZSbdFxGdN9jLRJH31Zbs1EfZdkmZP+P07kj5qoI9JRcRH1eMeSb/V+NuOQbL74Ai61eOehvv5m4jYHRFfRcTXkn6pBrddNcz4M5J+HRHPVpMb33aT9dWv7dZE2N+QNM/22bZPlPQjSWsb6OMwtk+pPjiR7VMk/UCDNxT1WklLq+dLJf2uwV6+YVCG8W41zLga3naND38eEX3/kXSVxj+R/x9J/9JEDy36mivp7ern3aZ7k7RG44d1/6fxI6KfSDpd0npJH1SPpw1Qb/+l8aG9N2k8WLMa6u0ijb813CRpY/VzVdPbrtBXX7Ybl8sCSXAFHZAEYQeSIOxAEoQdSIKwA0kQdiAJwg4k8f/6+GX5HZ2APwAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(xb[1].view(28, 28))\n",
    "yb[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 349,
   "id": "a70720b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "model, opt = get_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 350,
   "id": "4ce0af34",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss: 2.29, accuracy: 0.14\n",
      "loss: 2.31, accuracy: 0.06\n",
      "loss: 2.30, accuracy: 0.12\n"
     ]
    }
   ],
   "source": [
    "fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 353,
   "id": "1cb87759",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.multiprocessing as mp\n",
    "from fastcore.basics import store_attr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 354,
   "id": "622de54b",
   "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",
       " tensor([0, 9, 2]))"
      ]
     },
     "execution_count": 354,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_ds[[1, 4, 5]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 355,
   "id": "1f3a33a5",
   "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",
       " tensor([0, 9, 2]))"
      ]
     },
     "execution_count": 355,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_ds.__getitem__([1, 4, 5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 359,
   "id": "8ada49f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.]]), tensor([1, 9, 1]))\n",
      "(tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "        [0., 0., 0.,  ..., 0., 0., 0.]]), tensor([2, 1]))\n"
     ]
    }
   ],
   "source": [
    "for ele in map(train_ds.__getitem__, ([3, 4, 200], [5, 6])):\n",
    "    print(ele)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 360,
   "id": "87326dba",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataLoader:\n",
    "    def __init__(self, ds, batches, n_workers=1, collate_f=collate):\n",
    "        fc.store_attr()\n",
    "    \n",
    "    def __iter__(self):\n",
    "        with mp.Pool(self.n_workers) as ex:\n",
    "            yield from ex.map(self.ds.__getitem__, iter(self.batches))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 361,
   "id": "a10de849",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dl = DataLoader(train_ds, batches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 362,
   "id": "fa92fd44",
   "metadata": {},
   "outputs": [],
   "source": [
    "it = iter(train_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 367,
   "id": "994e3cb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "xb, yb = next(it)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 368,
   "id": "7a66d54d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([4, 784]), torch.Size([4]))"
      ]
     },
     "execution_count": 368,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xb.shape, yb.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 381,
   "id": "c99e7772",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8"
      ]
     },
     "execution_count": 381,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mp.cpu_count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0527ddc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "feca38e6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5b200a4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "701ec5bb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a49194d1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:python_main] *",
   "language": "python",
   "name": "conda-env-python_main-py"
  },
  "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
}