{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4f498b2e",
   "metadata": {},
   "source": [
    "# Generating the Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "37906477",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "# Loading visualization functions\n",
    "from visualize_sample import visualize_sample\n",
    "\n",
    "from misc import evaluate\n",
    "\n",
    "# Loading the NN model\n",
    "from model import NN_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60167a1d",
   "metadata": {},
   "source": [
    "Now we load the dataset and convert it to black and white."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9302f6c5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAB7sAAAGKCAYAAAB9+iVVAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy80BEi2AAAACXBIWXMAAA9hAAAPYQGoP6dpAAA3iklEQVR4nO3dCbgU1Z0/7mq8IooKLrizmBg1ahQmkiCKYswYnUQHN1DHqBg1MUGjaNwNYpRoIiqMkbiCTNSAC65RGR1RUG7UGIwOwQVBCRoDAi6gglC/p9t/MuN/Tl0p6L7dp/t9n8eH5FNddc6tW6frVH1vdRfSNE0TAAAAAAAAAIhIm2p3AAAAAAAAAADyUuwGAAAAAAAAIDqK3QAAAAAAAABER7EbAAAAAAAAgOgodgMAAAAAAAAQHcVuAAAAAAAAAKKj2A0AAAAAAABAdBS7AQAAAAAAAIiOYjcAAAAAAAAA0VHsBgAAAAAAACA6it014JJLLkkKhUKy0047VbsrUFV/+MMfkv322y9Zf/31k/XWWy/Zd999k2nTpvmt0LAmTZpUOj+E/mtubq5296BmmEvBpz744INkyJAhpfnUhhtuWDpfjBkzxu6hoZlPwcoxn4JPvfLKK8nhhx+ebLXVVsk666yTbL/99slFF12ULFmyxC6iIZlLQcuee+655MADDyxdgxfPG8U638iRI+22VtbU2g3yWX/5y1+SYcOGJe3bt7drSBr9pLDHHnsknTt3Lt2kXbFiRXLNNdcke+21V/L0008n2223XbW7CFVzyimnJD179vxMts0221StP1BLzKXgf8yfP790M7ZLly7JLrvsUroxBXzKfArMp+DzzJkzJ/na176WdOjQIRk0aFCpcDF16tTSfariAxr33HOPnUjDMpeC/2vixInJAQcckPTo0SO54IILknXXXTeZOXNm6V4VrUuxu8rOOOOMpFevXsny5ctLN6egURVPBmuvvXbpImKjjTYqZUcddVSy7bbbJueee25y5513VruLUDV9+vRJDj30UL8BCDCXgv+x+eabJ2+99Vay2WabJc8+++z/+UMpaGTmU5DNfAo+9R//8R/JokWLkilTpiQ77rhjKTvxxBNLD2SMHTs2WbhwYbLBBhvYXTQkcyn4rPfeey85+uijk29/+9vJHXfckbRp44O0q8ner6InnniiNAiuuuqqanYDasLkyZOTb37zm/8odP/9hm3xye7777+/9LGc0Mjef//95JNPPql2N6CmmEvBZ6211lqlQjcQZj4F5lPweYWLok033fQzefH+VLGI0bZtWzuQhmYuBf/j1ltvTd5+++3SV8EUzxGLFy8u/XEU1aHYXSXFJ7lPPvnk5Pjjj0++8pWvVKsbUDM+/vjj0pPd/3/F77lYunRp8uKLL1alX1ALBg4cWPou+3bt2iV777136Wk9aHTmUgDkYT4F5lPwefr27Vv693vf+14ybdq00seajxs3Lhk1alTpI5x9DSWNzFwKPuuRRx4p3a+dO3du6StYix9hXvz/J510UvLRRx/ZXa3Mx5hXya9//evk9ddfLw0IICmdEJqbm0vFizXWWKO0S4pF7t///vel/108aUCjKf7V+CGHHJL8y7/8S7Lxxhsn06dPTy6//PLSR0c99dRTpe+DgUZlLgXAyjCfAvMpWFn77bdf8rOf/SwZNmxYcu+99/4jP++885KLL77YjqQhmUtB2CuvvFL6FM5//dd/Lf2R1M9//vNk0qRJyb//+7+XvhLjtttus+takWJ3FbzzzjvJT3/609J3FHfq1KkaXYCa88Mf/rD0V0/FE8OZZ55Z+siP4oVE8Xsniz788MNqdxFaXe/evUv//d2BBx5Y+u7unXfeOTnnnHOShx56yG+FhmQuBcDKMp8C8ynIo1u3bsmee+5Z+sPz4lftPfDAA6Xid/GrYgYNGmRn0nDMpSCs+LWrS5YsSX7wgx8kI0eOLGUHH3xw6QG+a6+9NrnooouSL33pS3ZfK/Ex5lVw/vnnJxtuuGHpY8yBTxVPCueee27puy523HHH0sf7z5w5s1T4Lip+DAiQJNtss03pLwYfe+yx0ichQCMylwJgdZhPgfkUhPz2t79NTjzxxOSGG25ITjjhhFLR4sYbb0yOOeaY5Kyzzir90S1gLgVFf/9K1iOOOOIzO+TII48s/Tt16lQ7qhUpdlfhow2uu+660ve8vPnmm8ns2bNL/xU/w3/ZsmWl/71gwYLW7hbUhEsuuSR5++23k8mTJyd/+tOfkmeeeab0hHfRtttuW+3uQc3o3Llz6a8EFy9eXO2uQKszlwKgHMynaGTmUxB2zTXXlL4ubKuttvpMXvyUteLTe3/84x/tOvj/mEvR6LbYYovSv5tuuuln8k022aT078KFC6vSr0al2N3Kit87XCzeFYvdW2+99T/+K34v8csvv1z638WPN4BGtcEGGyR77LFH6cnuouL32hcvMrbffvtqdw1qxmuvvZa0a9fOJx7QkMylACgH8ykamfkUhBUfwAh9glrxAaWi4nezAp8yl6LRffWrX/3HvOp/Kz7kWuQrjFuX7+xuZTvttFMyYcKE4Mdxvv/++8mIESOSL37xi63dLahJ48aNKz3dffnllydt2vjbHBrPvHnz/s/E6Pnnn0/uvffeZP/99zcuaEjmUgDkYT4F5lOwsoqfKjhx4sTSA0n/+xMGb7vtttL1984772xn0nDMpSCsf//+yaWXXlr6uotvfOMb/8iLX4XR1NSU9O3b165rRYU0TdPWbJCw4oE/f/785MUXX7SLaEhPPPFE6VMN9t1332SjjTZKmpubk9GjRyf//M//nNx3332lEwQ0muJEqfj9L7179y59BM706dNLX4Wx5pprlr735ctf/nK1uwg1w1wKPnX11VcnixYtKv01+ahRo0rfNVn8OM6ik08+OenQoYNdRUMxn4KVZz5FoyvemyqeN4r3pQYNGlT69/77708efPDB5Pjjj0+uv/76ancRWp25FGT73ve+l9x0002lwvdee+2VTJo0Kbn99tuTc845Jxk2bJhd14oUu2uECwoa3cyZM5Mf/vCHyXPPPVf6lIPiR/ofc8wxyeDBg5O2bdtWu3tQFSNHjkxuueWW5NVXX03ee++90lPe++yzTzJkyJBkm2228VuB/8VcCj7VrVu35PXXXw/ujlmzZpWWQyMxn4KVZz4FSfL0008nF154Yen7ud95551/3J8688wzPYhBQzKXgmzFr7koFrWLD+0V/+C8a9euyY9+9KPk1FNPtdtamWI3AAAAAAAAANHxJbgAAAAAAAAAREexGwAAAAAAAIDoKHYDAAAAAAAAEB3FbgAAAAAAAACio9gNAAAAAAAAQHQUuwEAAAAAAACIjmI3AAAAAAAAANFpWtkXFgqFyvYEyixN01bZp8YGsWmNsWFcEBvnDKje2HDOIDbGBVRnXBQ5ZxAbYwOqNzacM4iNcwas+tjwZDcAAAAAAAAA0VHsBgAAAAAAACA6it0AAAAAAAAAREexGwAAAAAAAIDoKHYDAAAAAAAAEB3FbgAAAAAAAACio9gNAAAAAAAAQHQUuwEAAAAAAACIjmI3AAAAAAAAANFR7AYAAAAAAAAgOordAAAAAAAAAERHsRsAAAAAAACA6Ch2AwAAAAAAABAdxW4AAAAAAAAAoqPYDQAAAAAAAEB0FLsBAAAAAAAAiI5iNwAAAAAAAADRUewGAAAAAAAAIDqK3QAAAAAAAABER7EbAAAAAAAAgOgodgMAAAAAAAAQHcVuAAAAAAAAAKKj2A0AAAAAAABAdBS7AQAAAAAAAIhOU7U7AAAAAAAAAFDLunXrFsz/9Kc/Za6z3nrrBfPu3bsH8+eff34Ve9e4PNkNAAAAAAAAQHQUuwEAAAAAAACIjmI3AAAAAAAAANFR7AYAAAAAAAAgOordAAAAAAAAAESnqdodAACAWP3qV7/KXHbooYcG8z59+gTzl19+uWz9ov4cfvjhwfy2224L5gceeGDmtu67776y9QsAKmXWrFnBfODAgcF80qRJfhkR69y5c+559WmnnZZrW1OnTg3mAwYMCOZz5szJ7BNUW9euXXO9fuONNy5Lu1/+8pczl82fPz+Y9+vXryxtn3jiiZnL0jQN5oVCIZjfddddud5v4O9uuOGG4M5Yd911cx+flI8nuwEAAAAAAACIjmI3AAAAAAAAANFR7AYAAAAAAAAgOordAAAAAAAAAERHsRsAAAAAAACA6DRVuwO1ol27dsF8nXXWCeZLliwJ5h999FFSLcOGDcvV11/+8peZ2/r444/L1i8oWnPNNYM7onPnzsF8yJAhmTvu6KOPruhO/cEPfpC57LrrrgvmaZpWsEcA1KoVK1ZkLuvUqVMw33rrrYP5yy+/XLZ+UX8WLlyY6xjcd999M7d13333la1fQOXss88+wXzdddcty/bvueeesmwHVsaFF14YzHfZZZfMdbp16xbM+/btG8wnTZrklxGB/v37B/Nx48ZlrjNnzpxgfsYZZ+Rq+9BDD83Vdu/evXNtH1rT2LFjg/l2222X6/o0655moVDI9fpVWSfv61tqO++92X79+uV6PY0n6xzQs2fP3Nv64IMPgvn777+fe1uEebIbAAAAAAAAgOgodgMAAAAAAAAQHcVuAAAAAAAAAKKj2A0AAAAAAABAdBS7AQAAAAAAAIhOU7U7UCuuu+66YH7UUUcF84kTJwbz/fbbL6m0XXfdNZifeeaZwbxNm/DfNLzzzjuZbYwaNWoVe0ej2HLLLYP5ySefHMwPOuigYL7NNtvkbnvFihVJJV1zzTWZy5YuXRrMR48eXcEeQfmtvfbawfzDDz/Mva311luvDD1KkjXWWCOYL1q0qCzbp/bccsstwfy2224L5vfff39Sa2Nm//33z1xn8eLFwfzdd98tW79oHM8//3yu1w8YMCBz2a9+9atgPmPGjNz9glp0xBFHBPPDDjssc52s+fwrr7ySq+0OHTrkuh4q6tOnTzDv2bNnrjnT3/72t2C+YMGCYH7PPfdk9gn+rmPHjsGdceWVVwbzY489tuI7b/bs2RVvg9XXq1evYD5u3LhgPnjw4MxtZR1veY0fPz5Xftppp1W8T/379891362cbVP7zjvvvNzzhzRNg3mhUMjVdt7Xt0Ybq9KnrNrIxRdfnHtbNJZLLrmkbPdBf/3rXwfz1157Lfe2CPNkNwAAAAAAAADRUewGAAAAAAAAIDqK3QAAAAAAAABER7EbAAAAAAAAgOgodgMAAAAAAAAQnaZqd6BWvP7668E8TdNgPnXq1KRa/umf/imYFwqFXNvZZJNNytQj6tmGG24YzCdOnBjMt99++6SenXbaacF85513DuaPP/54ML/77rvL2i/i1K1bt2DesWPHYD5t2rTcbVx44YXB/Dvf+U4wnzlzZu42evTokescmmX99dcP5k888UQwHzBgQK7tUz1du3YN5gcffHAwf/bZZ5Nac8EFFwTzL3zhC5nrPP/888G8ubm5bP2CLBtttFHmsl133TWYz5gxww6lJmUds1deeWUw79mzZzBfc801M9vo169fWeYzq2LKlCnB/Oabbw7mjz76aDCfPHlyMJ87d+5q9I5G0Ldv38xlo0ePznUtk2XMmDG5xl5L10Wup+PQpUuXXK+/4447kmq54oorgvn48eNz93fOnDm57idltT148ODMtmkcLb1HZs1R8s5d8r5++vTpuec0eV//5z//Oak01z583n2dr371q7l20uLFizOXXXPNNXZ4hXmyGwAAAAAAAIDoKHYDAAAAAAAAEB3FbgAAAAAAAACio9gNAAAAAAAAQHQUuwEAAAAAAACIjmI3AAAAAAAAANFpShrICSeckLnsjDPOCOaffPJJMB86dGhSaSeddFIwv+qqq4J5oVDItf2XX355lfpF/VljjTUylx1zzDHBfPvtty9L27NmzQrmw4cPT8pll112yf2ekGXHHXfMla9YsSKY33333bnbJl79+vUL5ldeeWUw/+ijj4L57rvvntnGnnvuGcwHDx4czNdZZ51g3r179ySvrPNPmqZJORxyyCFl2Q7Vc8ABBwTztdZaK4lF1lwRWss777wTzO+///5g/p3vfCdzW4ceemgw/81vfrOKvYPK2mqrrXLPjUL+9Kc/ZS578MEHy3Ld/N577wXzO++8M9d2oBJOPfXUYD5kyJDc2xozZkwwP+2003JdZxx77LGZbSxatChXTm0ZP358rnnIuHHjMrfVu3fvpJKam5vLdg3Qv3//YH7FFVfkyrPuF1CfunbtGsy7dOmSuU7eWsCIESOC+bBhw4L5/Pnzc20fYjVq1Khgvu666+bazrXXXpu5bPbs2Uk5dOrUKZjPmzcvaXSe7AYAAAAAAAAgOordAAAAAAAAAERHsRsAAAAAAACA6Ch2AwAAAAAAABAdxW4AAAAAAAAAotOUNJAf/ehHmcvWWmutYH7ttddWsEdJstVWW2UuO/fcc4P5mmuumauNv/71r8H8oYceyrUd6teXvvSlzGWXX355WdrIOg6vuuqqYP7ggw/mbmObbbYJ5gcddFDubcGq2GOPPYL5nXfemWs7r7zySu62v/rVrwbz9u3bJ5W2dOnSYP7qq68G82eeeSZXTjy23XbbYP6Tn/wkmK9YsSKYT5s2LamW7t27B/M2bfL/jeioUaPK0CP41LJly4K7YuHChXYRdeP73/9+MB86dGgwT9M0mD/88MPB/KSTTsps+/XXX1+pPkIt6datWzAfPXp0MO/bt28wv/vuuzPbKNf1dMeOHXOv01K/iNcVV1wRzMePH5+5zlNPPRXMe/funVRSS33q3LlzMD/11FOD+YABA3K3QePo06dPMN9oo40y18maB2XlgwcPXsXeQX1Yf/31g3mPHj1ybWfx4sXBfOTIkbn7dPjhhwfz8847L5hvscUWue+pDsg4/7z77rtJPfFkNwAAAAAAAADRUewGAAAAAAAAIDqK3QAAAAAAAABER7EbAAAAAAAAgOgodgMAAAAAAAAQnaZqd6BWvPDCC8H8ggsuKMv2t9hii2B+3333Za6z5ZZblqXtZcuWBfOFCxeWZfvE78gjj6x4G5tttlkwHzFiRK4catmAAQOCeZqmubYzd+7cYL5gwYLMdS699NJg/sgjjySVtnjx4mD+7LPPVrxtWl+bNtl/K9mrV69g3rlz52D+6KOPBvPHHnssqbRddtklmA8bNizXz/36669ntnHTTTetYu8A6ts3v/nNYD5q1KiyzKWefvrpYN69e/fMdbLOVVOmTMnVNlRC3759g/mECRNybWfgwIHBfMyYMUml7bXXXrnXufnmmyvSF6qrubk5mO++++6Z6zz55JPB/KmnngrmvXv3Tipt+PDhwXzq1KnBfPz48RXuETE799xzg3mhUMhcp6VlIY8//ngw//Of/5yUS9a8afLkybmvp6Hcjj766GC+8cYb59rO73//+2D+xhtvZK5z/PHHB/Orr746mLdt2zZXn/bdd9/MZT169AjmkyZNSuqJJ7sBAAAAAAAAiI5iNwAAAAAAAADRUewGAAAAAAAAIDqK3QAAAAAAAABER7EbAAAAAAAAgOgU0jRNV+qFhUISu2nTpmUumzx5cjA/+eSTc+2P448/PpiPGDEimLdr1y6zT3/84x+D+f333x/ML7jggmD+29/+NpgfeeSRST1byUN7tdXD2OjUqVPmsj//+c/BfIMNNqhgj+Lz0ksvBfPjjjsumDc3Nyf1PDbqYVysimOPPTaY33DDDbm2s2zZsmB+7bXXZq4zbNiwYL5o0aJgvnTp0lx9qnfOGf9X27Ztg/vqF7/4ReZ+POWUU4L54sWLg/mmm24azJcsWZJU2tixY4P5UUcdlesYyZr7FY0ePTqJnXNG7Tv66KNzH38TJ04M5gcccEAw/+STT1axd/XJuFh9O+20UzAfP358MG/TJvx3+l/60peScsk6zqdMmRLML7300mD+X//1X8F8+fLlST0zl1p5HTt2zFw2ZMiQYH7qqafmus81cODAXK9vjZ9v1qxZwXz27NmZ2+rRo0cSO2OjPDp37hzMn3zyyVz3XPr375+r3dNOOy1z2WGHHRbMe/funauNRmU+tXLzhJb2U9a9t6x1Kv36ltaZN29eMD/ppJOC+YQJE5JG5JxRWTNnzgzmW2+9da7tdO/ePZgffPDBmeucffbZue69ldMhhxwS/ThbmbHhyW4AAAAAAAAAoqPYDQAAAAAAAEB0FLsBAAAAAAAAiI5iNwAAAAAAAADRUewGAAAAAAAAIDpN1e5ArVhnnXWC+ZZbbhnMf/rTnwbzE044IZgvW7YsmA8fPjyzTz//+c+D+THHHBPMly5dGswvu+yyzDagaN68eZk7YuTIkcH85JNPDuYbbrhh9Dt1wYIFmcvOOOOMYD5u3Lhg/tFHH5WtX9S+u+66K5h/8MEHwfz6668P5uutt14wHzRoUGbbWWPyxhtvzLWtrHMJjWefffYJ5qecckrmOu+9914wP/TQQ4P5kiVLkkq74IILgvmBBx6Yazvjx48P5g899FDmOl/4wheC+dtvvx3Mu3btGswXLVoUzN98883Mtmkczz77bO519t1332C+zTbbBPMZM2bkbgNa8uKLLwbzHXbYIZivueaawbxnz57BvEOHDsH84IMPzuxTnz59gnnfvn1z5cOGDct1PqJ+dezYMZiPHj06c51+/foF86uuuiqYDx06NNfcoTVkjY2s/TFixIgK94h6MGfOnGC+++67B/Mnn3wy15z+iiuuyJUXDR48OHMZ5FUoFCq+TqVf39I6nTp1CuZ33nlnME/TNNe9haIJEyasVB+pf+uuu26ua4Qsc+fODeZbbLFFMD/77LMzt9W2bdtgPmvWrFw1iKwx05Ks/tYbT3YDAAAAAAAAEB3FbgAAAAAAAACio9gNAAAAAAAAQHQUuwEAAAAAAACIjmI3AAAAAAAAANFpqnYHasXAgQNz5Vnmzp0bzC+77LJgfvXVVyd5bbjhhsF82bJlwfz555/P3Qb83UUXXRTcGZMnTw7mF198ca6dt/322wfzjh07VvyXcOuttwbzyy+/PHMd44mWvPfee8H8jjvuCOZ77LFHMB80aFDZdvRxxx0XzM8777xgPm/evLK1TRx23XXXYH7jjTfm3tYDDzwQzBctWhTM11hjjWC+fPnypFz23XffYN6hQ4dgnqZpMO/evXswf+qppzLbzpqbPfPMM8H8iCOOCObXX399MP/+97+f2TZAPcl6P23pPTjkwQcfzFy26aabBvOzzjormB999NG55ljPPfdcZtsTJkzIXEbt69atW67fa9acomjo0KHB/MILL0xiccwxx+R6/ezZsyvWF+rfnDlzgnmXLl2C+fjx44P51KlTc+VFV1555Ur1EVbG4MGDg/l2222Xuc6JJ54YzKdPnx7Mp0yZUvFfRlafsq6zs2S9fuzYsZnr9OzZM5jPmDEjV9vEL+tYyKqrZVm4cGEw/+UvfxnM27Ztm7mtpUuX5roPdOSRRybl8t///d9JI/BkNwAAAAAAAADRUewGAAAAAAAAIDqK3QAAAAAAAABER7EbAAAAAAAAgOgodgMAAAAAAAAQHcVuAAAAAAAAAKJTSNM0XakXFgpJ7E466aTMZVdeeWUwb9u2bTCfNWtWMB80aFAwf/DBB5O8Nt1002A+ffr0YH7vvfcG84EDByaNaCUP7dVWD2OjNay33nrB/OGHHw7mX//618vW9nPPPRfMDzvssGA+e/bspJ61xtgwLlZP1vv2xRdfnLnOFltsEcxXrFgRzDfbbLNgPm/evKQRNfI547rrrgvmxx9/fO5tzZ8/P5ivv/76wfyPf/xjMP/kk0+Sctl5551znZdaQ3NzczDv0KFDMB81alQwv/rqq5NKc86ofTvssEMwf+GFF3Jva8cddwzmM2bMyL2temZcUPTd7343uCPGjBkTzG+66abMHXfCCSdEv1MbYS7VsWPHXPeHsl4/dOjQzDYuvPDCJHZZ+2PRokXBvEePHkk9a4SxUc+/pzlz5mSu079//1xzfVZun5eTcdH6OnXqFMzPOeecYH7qqafmOj5a+p3uv//+ue491yLnjPLYe++9g/mjjz6aVMsBBxyQq/44duzYYN6+fftgvnz58sy2+/TpE/35amXGhie7AQAAAAAAAIiOYjcAAAAAAAAA0VHsBgAAAAAAACA6it0AAAAAAAAAREexGwAAAAAAAIDoNCUNZNSoUZnLnnnmmWC+1lprBfNZs2YF8zfffDMpl9/85jfBfIMNNgjmDz/8cNnahnL72c9+Fsy//vWvl62N5ubmYP6tb30rmH/wwQdlaxvKafTo0cF84sSJmeu88cYbwTxN07L1Cz7PxhtvnGsnlescUCgUMpdljYGsdbJef+ONNwbz4cOHJ3n95S9/yTXvfOedd3K3QeOYOXNmMJ80aVLmOn379g3mvXr1CuYzZsxYxd5B/erYsWO1u0Ar/24fe+yxXNsZOHBgMB8zZkwS0/HcvXv3XOeSbt26BfPZs2cH8wsvvDCz7a5du+ZqI0vW67feeutc26HxdO7cOXPZ+PHjg/kZZ5yR6/VQT+bNmxfMBw8eHMx//OMfl+1eVr9+/YK5mknj2X777avS7rvvvpu57Kijjgrm3/nOd4J5+/btc7U9dOjQ3DWTeuPJbgAAAAAAAACio9gNAAAAAAAAQHQUuwEAAAAAAACIjmI3AAAAAAAAANFR7AYAAAAAAAAgOk3V7kCtePbZZ6vSbufOnTOX7bDDDrm2deedd5ahR7By2rdvH8wPO+ywYH7UUUeVbdfOnz8/mJ9//vnB/IMPPihb29S+Rx99NJiPGzcumF933XVJrWnbtm0w/+lPf5q5zvLly3Pl8HeXXHJJcGe8//77Fd9JvXr1Cua77bZbru2kaZq57LXXXsv1c//ud78L5n/7299yt52X8xWr4uOPPy7b8eScQUu+8IUv5Jq3HHHEEZnbWmuttYJ5oVCo+Httlqy2v/KVrwTzvffeO5gvXbo0mE+ePHk1ekdrGj16dDDv3r17MB8zZkwwnz17djDv27dvZtvdunXLlXfo0CFXX1tqu1r22muv3Os8/vjjwXzatGm5cvg8t99+e+ayqVOn5rr3cOihhwbz008/PZjPmTPHL4i60alTp1zzryx5X09jateuXVXazZqXFQ0YMKAs1+Y33HBDMB8xYkTS6DzZDQAAAAAAAEB0FLsBAAAAAAAAiI5iNwAAAAAAAADRUewGAAAAAAAAIDqK3QAAAAAAAABEp6naHWh0PXr0yFy2+eabB/Nrr702mH/yySdl6xd8nm984xvB/MYbb6z4zhs0aFAwf+yxxyreNrWvb9++wfzee+/N9fpVscEGGwTzdu3aBfNvfetbwbxLly7BfM8998xs+6yzzgrmixYtCubLli3L3BaN5fXXXw/mp59+esXbXmONNYJ5//79g/ktt9ySu41jjz02mE+ZMiX3tqDeNTW5PCRJJk+eHNwN3bt3D+brrLNO2XZboVAI5mmaVvxXU662f/CDHwTzsWPHrlK/qH1Zc42svCVZc/dp06bl2s7jjz8ezG+++eZgPnv27MxtZY39K6+8MphPmjQpmO+9996ZbUC1DRgwIJiPGzcuc52s66W5c+cG88svvzxXG1l9KpozZ07mMqhFBx10UK55VmvM/ahfzzzzTDBfsWJFMG/TpnrPBC9YsCCYX3bZZcH8l7/8ZYV7FC9PdgMAAAAAAAAQHcVuAAAAAAAAAKKj2A0AAAAAAABAdBS7AQAAAAAAAIiOYjcAAAAAAAAA0Wmqdgca3R577JF7nSVLlgTzNE3L0CP4rK997WvBXXLGGWdUdFc9++yzmct+97vfVbRt6tPw4cMr3kahUKja+3Nr/HxQbitWrAjm++yzT67tTJs2LXPZzJkzc/cLGtXUqVOr3QVq+Bq1NeYzWXOp1pDV9pNPPhnMhw4dGswfeeSRsvaL1nfQQQcF8+7duwfzjh07lm3esmjRoqTW9O3bN9frs8YG1NscKOsavH///rnayDrPtHSNn9UG8eratWsw33jjjTPX+cMf/pDEok+fPmWZ+7X0+uuvvz53v6hPU6ZMCeZvvfVWMN9yyy3L0u7y5cszl91www3B/Cc/+Ukw/+CDD8rSp0biyW4AAAAAAAAAoqPYDQAAAAAAAEB0FLsBAAAAAAAAiI5iNwAAAAAAAADRUewGAAAAAAAAIDqK3QAAAAAAAABEp6naHWgUX/ziF4P5UUcdlblOoVAI5o888kjZ+gVFm2yySeaOePjhh4P5+uuvX9Gd19zcnLls8eLFFW2buGW9d8bU9pIlS4L5/fffX5btQ6344Q9/GMyPO+64YD537txgfsABB2S28dZbb61i7wAa0+jRo4P5t7/97Vzb2XDDDTOXtWmT7+/u0zQN5u+9914wX7p0aTB/4IEHMtu46667gvnEiRNztUH9mjZtWlKvunfvnrlsyJAhwXzMmDHBfNKkSWXrF7SWOXPmBPOpU6dmrtOrV6+ytNG/f//cbQ8fPjyYn3766bn6RO14+umng/lGG22Uuc5hhx0WzCdMmJBUy3nnnRfM+/Xrl2uOl5XPnz8/s+2WlkHRwoULgztiyy23zHU/6ZZbbsl1PfF5tQ7Kw5PdAAAAAAAAAERHsRsAAAAAAACA6Ch2AwAAAAAAABAdxW4AAAAAAAAAoqPYDQAAAAAAAEB0mqrdgUax+eabB/PNNtssc500TYP5c889V7Z+0VjatWsXzK+99trMddZff/0K9ihJbrrppmB+1llnVbRd6tfee++d6/U9e/YM5vvvv3/mOh9//HEw32233YL5q6++mqtP9913XzC/6KKLcm0Hat2ll16a6/UvvPBCMJ87d26ZegSNrWvXrsF8xowZrd4Xque4444ry3b222+/zGVt27YtSxtPP/10MP/rX/9alu1Do5gwYULmskWLFgXzoUOHVrBHUBsGDx6cuWzq1KnBfPjw4cH89NNPD+bNzc252z7ttNNytUHtOPHEE4P5JptsEsxXrFiRua1zzz0393t6yJ577hnM/+3f/i2Y9+vXL3NbnTp1ylXnKBQKSR4jR47MXPbGG2/k2haNZ9NNN831+h//+MfB/I477gjmXbp0WaV+UR6e7AYAAAAAAAAgOordAAAAAAAAAERHsRsAAAAAAACA6Ch2AwAAAAAAABAdxW4AAAAAAAAAotNU7Q4ArWevvfYK5gceeGDF237hhReC+aBBg4L5xx9/XOEeUa8ef/zxsrx+5MiRmeusWLEimK+33nrB/KOPPsrVpw8//DDX66HW9e/fP5ivvfbauc4Bv/jFL8raL+CzBgwYENwlDz/8sF1Fbg899JC9BjWmW7duufKiq666KpjPnj27bP2CWtXc3Jx73nT55ZcH89122y2YDx48OJjPnTs3s+3OnTsH8169euX+OagNWfeZ0jTNXGfjjTcO5ieeeGIw79SpUzA//vjjg3mXLl1y9ylrWUvrhEyfPj2YX3LJJbm2A6vj6aefzvX6N954ww6vIk92AwAAAAAAABAdxW4AAAAAAAAAoqPYDQAAAAAAAEB0FLsBAAAAAAAAiI5iNwAAAAAAAADRaap2B8hvu+22C+YdO3YM5i+99JLd3GDatm0bzM8+++ykWjbYYINgvvnmmwfz2bNnV7hH0LKlS5fm3kULFy60WyHg/PPPD+6XNm3Cf3d52WWXBfNJkybZv7CS/vM//zNzWc+ePYP5V77ylWDevn37YL548WK/D4CIDBkyJPc6I0aMqEhfIHbjx48P5m+88Uau10+dOjVXXjRnzpxg3tzcnLkOteGJJ57Idf8+qw5Q1K1bt2A+atSoYF4oFIJ5mqa5Xt+SvOvceuutwfy73/1u7rbh82S9Py9YsCCYz58/306NiCe7AQAAAAAAAIiOYjcAAAAAAAAA0VHsBgAAAAAAACA6it0AAAAAAAAAREexGwAAAAAAAIDoNFW7A+R3zz33BPPvfve7wfyll16ymxtMoVAI5u3bt0+qZc011wzm7dq1a/W+ANC6xo4dG8wvvfTSYP7II49UuEdQ/66++urMZd26dQvmr776ajBfvHhx2foFQOVlvc937NgxmF911VWZ25o9e3bZ+gWNoLm5OZjvvvvuwXzcuHHBfLfddstsY/DgwavYO6ptxowZwXzXXXfNdS1ddNBBBwXzNE1z9Snv66dPn565bMqUKcF8woQJwXzixIm52obV0bNnTzuwjnmyGwAAAAAAAIDoKHYDAAAAAAAAEB3FbgAAAAAAAACio9gNAAAAAAAAQHQUuwEAAAAAAACITiFN03SlXlgoVL43daypqSmYd+/ePXOdhx56KJi/9tprwXy33XYL5suXL08a0Uoe2qstprGxySabBPO33nor97YWLFgQzK+55ppgfvvttwfzF198MXfb1P7YiGlcQJFzBoQ5Z4BxASvDXAqMDcjDdQaYT0E5zxme7AYAAAAAAAAgOordAAAAAAAAAERHsRsAAAAAAACA6Ch2AwAAAAAAABAdxW4AAAAAAAAAoqPYDQAAAAAAAEB0Cmmapiv1wkKh8r2BMlrJQ3u1GRvEpjXGhnFBbJwzoHpjwzmD2BgXUJ1xUeScQWyMDaje2HDOIDbOGbDqY8OT3QAAAAAAAABER7EbAAAAAAAAgOgodgMAAAAAAAAQHcVuAAAAAAAAAKKj2A0AAAAAAABAdBS7AQAAAAAAAIiOYjcAAAAAAAAA0VHsBgAAAAAAACA6it0AAAAAAAAAREexGwAAAAAAAIDoKHYDAAAAAAAAEJ1CmqZptTsBAAAAAAAAAHl4shsAAAAAAACA6Ch2AwAAAAAAABAdxW4AAAAAAAAAoqPYDQAAAAAAAEB0FLsBAAAAAAAAiE7Tyr6wUChUtidQZmmatso+NTaITWuMDeOC2DhnQPXGhnMGsTEuoDrjosg5g9gYG1C9seGcQWycM2DVx4YnuwEAAAAAAACIjmI3AAAAAAAAANFR7AYAAAAAAAAgOordAAAAAAAAAERHsRsAAAAAAACA6Ch2AwAAAAAAABAdxW4AAAAAAAAAoqPYDQAAAAAAAEB0FLsBAAAAAAAAiI5iNwAAAAAAAADRUewGAAAAAAAAIDqK3QAAAAAAAABER7EbAAAAAAAAgOgodgMAAAAAAAAQHcVuAAAAAAAAAKKj2A0AAAAAAABAdBS7AQAAAAAAAIiOYjcAAAAAAAAA0VHsBgAAAAAAACA6it0AAAAAAAAAREexGwAAAAAAAIDoKHYDAAAAAAAAEB3FbgAAAAAAAACio9gNAAAAAAAAQHSaqt0BAAAAAAAAgBilaZp7nUKhUJG+NCJPdgMAAAAAAAAQHcVuAAAAAAAAAKKj2A0AAAAAAABAdBS7AQAAAAAAAIiOYjcAAAAAAAAA0WmqdgcAACBWaZrmXqdQKFSkLzSmrGPQcQZA7JzjWNVjJC/zJohjrLYG7wfU0/HcSDzZDQAAAAAAAEB0FLsBAAAAAAAAiI5iNwAAAAAAAADRUewGAAAAAAAAIDqK3QAAAAAAAABEp6naHaB80jQN5oVCwW6mbo7n1mDMAAD1PJcy1wEg5mt897/q06ocI3nnNFltOKaIUTXvnQLUGk92AwAAAAAAABAdxW4AAAAAAAAAoqPYDQAAAAAAAEB0FLsBAAAAAAAAiI5iNwAAAAAAAADRaap2B2pFmqZl2U6hUEhi+dlqsa/Er1xjqVZ/BuMGoPLvt95rAai3eX4tXic53xLrMe3YjVs1rwGy2qjF92gAqsu5IS6e7AYAAAAAAAAgOordAAAAAAAAAERHsRsAAAAAAACA6Ch2AwAAAAAAABAdxW4AAAAAAAAAotNU7Q7EqlAoVLsLUFVpmjbkbyDvz+29gmqr57FqfNGa6nksAcSqnO/NMb3PmwPRmmIaG1Du99WWjv9yvRdnteG9ns87BuuB45xG4VivPE92AwAAAAAAABAdxW4AAAAAAAAAoqPYDQAAAAAAAEB0FLsBAAAAAAAAiI5iNwAAAAAAAADRUewGAAAAAAAAIDpN1e4A+aVpardR9wqFQl2Mmay2y/nzEa+8x6bjBoBKnn+cZ8C4oDGV85o577nEPS5W5lio1hylXu5NQUtcA9DovD/XB092AwAAAAAAABAdxW4AAAAAAAAAoqPYDQAAAAAAAEB0FLsBAAAAAAAAiI5iNwAAAAAAAADRaUoaSJqmSUxi6y/1qdLHYaFQqOj2oVHez+vhnOH9oH7Vw/G5KhzTALV1HvG+TKMo59zLuKHcso6plo7brGW1eHzG1Fcai2MQ4hlj5nL5ebIbAAAAAAAAgOgodgMAAAAAAAAQHcVuAAAAAAAAAKKj2A0AAAAAAABAdBS7AQAAAAAAAIhOU7U7UOsKhUJFt5+maUW3D7V+HBoD1JNqHs+VPl9Bo3BeAvAeDPU6b6nmNcOqzLFc4zSWln7fWcdPVt4ax47rBmI7nlqjDe/bkI9zSfl4shsAAAAAAACA6Ch2AwAAAAAAABAdxW4AAAAAAAAAoqPYDQAAAAAAAEB0FLsBAAAAAAAAiE5TtTvQKNI0rXYXIFOhUMhc5tit3L6lcX7f5RpHLW3HsUZrKed5oZrHbbl+DmOPej/3QD2pl3MY9cv8BMoz38nKW+O92/mBRpb3PGa8UE/zLNfg1eXJbgAAAAAAAACio9gNAAAAAAAAQHQUuwEAAAAAAACIjmI3AAAAAAAAANFR7AYAAAAAAAAgOk3V7kCs0jSteBuFQqEsfcq7HViZ46c1xkAtMp6o1XGUd1uOZVqT4w2AWjhfmHtRC8p5HNbDHKsefgbiudbOe++0pfHq2IXV514W5DvHNGpNZmV4shsAAAAAAACA6Ch2AwAAAAAAABAdxW4AAAAAAAAAoqPYDQAAAAAAAEB0FLsBAAAAAAAAiE5TtTtQ69I0Lct2CoVCUmt9gkoc0/VwfJZzvALEqh7ez1vj56jmfnK+Aqj8+2a9nA9pfeU8dmI65xszxHjPynFLtdXzfdZyaml/xHSuJG6tMccz9vPzZDcAAAAAAAAA0VHsBgAAAAAAACA6it0AAAAAAAAAREexGwAAAAAAAIDoKHYDAAAAAAAAEB3FbgAAAAAAAACi05Q0kEKhkLksTdOKtwGNIO8YKNfYWxXGK9UW03ihfrXGceXYXT3OVwDxz9daOhd6n49bOec5jgWorKwx5nqFWlUv5wVjjEZXL2O5lnmyGwAAAAAAAIDoKHYDAAAAAAAAEB3FbgAAAAAAAACio9gNAAAAAAAAQHQUuwEAAAAAAACITlO1O1ArCoVCUmvSNI3+Z4BVPZ7LydigXrR0LFdzjEEjc44BgMbgHs3q7Y9y7ddyMo+jEseu4wpWflyU8xxgTFLLKj3fKagNerIbAAAAAAAAgPj4GHMAAAAAAAAAoqPYDQAAAAAAAEB0FLsBAAAAAAAAiI5iNwAAAAAAAADRaap2BwDKpVAo2JkkaZrW7fGR9bPB6sgaG/V+vNXDewJALav380hezjvxqMVjtxb7VIuMM2rheMs7Xh23ALD6PNkNAAAAAAAAQHQUuwEAAAAAAACIjmI3AAAAAAAAANFR7AYAAAAAAAAgOordAAAAAAAAAESnqdodAOpLmqYVb6NQKFS8DYhtvBgXlFstHlOrMmZq8ecAaPT5fD1z3qGWx1I5j8+8P5+xQS3LOj7Lef2Rta2s3JgBaFzOAfl5shsAAAAAAACA6Ch2AwAAAAAAABAdxW4AAAAAAAAAoqPYDQAAAAAAAEB0FLsBAAAAAAAAiE5TtTvQ6NI0rXYXYJU4domNYxbiYbwCUOsKhUK1u0CF+N2u3rzM/qPRx0bWGMjKs7bT0tgzzohNa1zjGxfEyHFbPp7sBgAAAAAAACA6it0AAAAAAAAAREexGwAAAAAAAIDoKHYDAAAAAAAAEB3FbgAAAAAAAACio9gNAAAAAAAAQHSaqt0B8isUCnYbrSZNU3sbgJpmbgSAcwkA9XQt09L9uKxlrosa6/6r3zfUFmOyujzZDQAAAAAAAEB0FLsBAAAAAAAAiI5iNwAAAAAAAADRUewGAAAAAAAAIDqK3QAAAAAAAABEp6naHQCqL03TpNYUCoVqd4FIlevYqea4cPzTKGrx/APQ6MxDoPGsypzMewWNoKXjPGvcZOV5x8yqtE1jKdexlnf71eTcg+ONLJ7sBgAAAAAAACA6it0AAAAAAAAAREexGwAAAAAAAIDoKHYDAAAAAAAAEB3FbgAAAAAAAACi01TtDgCNrVAoVLsLEOTYhNpjXEJlpWlq7AEArMa1SdZ8yjyLWpV1bALExJPdAAAAAAAAAERHsRsAAAAAAACA6Ch2AwAAAAAAABAdxW4AAAAAAAAAoqPYDQAAAAAAAEB0mqrdAaD1pGlqdwPQ6px/oPUVCoXMZcYkQGNalff/ls4nwMqPmazxZ15GudXDMeXcA+TlyW4AAAAAAAAAoqPYDQAAAAAAAEB0FLsBAAAAAAAAiI5iNwAAAAAAAADRUewGAAAAAAAAIDpN1e4A+aVpGswLhYLdCQDUnKw5ijkN1BbXEwAAtXFNtCrbIl7lPD5qkWOWWuA4rG+e7AYAAAAAAAAgOordAAAAAAAAAERHsRsAAAAAAACA6Ch2AwAAAAAAABAdxW4AAAAAAAAAotNU7Q40ukKhkLksTdNW7QuNe7y1xrHW0rEOQGNybgBjDwBzL2hkrolwfACsPk92AwAAAAAAABAdxW4AAAAAAAAAoqPYDQAAAAAAAEB0FLsBAAAAAAAAiI5iNwAAAAAAAADRUewGAAAAAAAAIDpN1e4A2QqFgt1Dq3CsAQAAAAAAEBtPdgMAAAAAAAAQHcVuAAAAAAAAAKKj2A0AAAAAAABAdBS7AQAAAAAAAIiOYjcAAAAAAAAA0VHsBgAAAAAAACA6it0AAAAAAAAAREexGwAAAAAAAIDoKHYDAAAAAAAAEB3FbgAAAAAAAACio9gNAAAAAAAAQHQKaZqm1e4EAAAAAAAAAOThyW4AAAAAAAAAoqPYDQAAAAAAAEB0FLsBAAAAAAAAiI5iNwAAAAAAAADRUewGAAAAAAAAIDqK3QAAAAAAAABER7EbAAAAAAAAgOgodgMAAAAAAAAQHcVuAAAAAAAAAJLY/D9mjKvRwQ7fPwAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 2000x400 with 20 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Load the original MNIST dataset\n",
    "mnist_original = datasets.MNIST(\n",
    "    root=\".\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=transforms.ToTensor()\n",
    ")\n",
    "\n",
    "# This is our training dataset\n",
    "# Convert grayscale to pure black and white (0 or 1)\n",
    "threshold = 0.5\n",
    "mnist_bw = datasets.MNIST(\n",
    "    root=\".\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Lambda(lambda x: (x > threshold).float())\n",
    "    ])\n",
    ")\n",
    "\n",
    "# This is our testing dataset\n",
    "mnist_bw_test = datasets.MNIST(\n",
    "    root=\".\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Lambda(lambda x: (x > threshold).float())\n",
    "    ])\n",
    ")\n",
    "\n",
    "visualize_sample(mnist_original, mnist_bw, num_samples=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "72bbd0c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using CPU\n",
      "Epoch 01 | Train loss: 0.6791 | Test loss: 0.3215 | Test acc: 90.22%\n",
      "Epoch 02 | Train loss: 0.2942 | Test loss: 0.2641 | Test acc: 91.88%\n",
      "Epoch 03 | Train loss: 0.2454 | Test loss: 0.2293 | Test acc: 93.13%\n",
      "Epoch 04 | Train loss: 0.2126 | Test loss: 0.2051 | Test acc: 93.82%\n",
      "Epoch 05 | Train loss: 0.1887 | Test loss: 0.1905 | Test acc: 94.08%\n",
      "Epoch 06 | Train loss: 0.1690 | Test loss: 0.1768 | Test acc: 94.53%\n",
      "Epoch 07 | Train loss: 0.1533 | Test loss: 0.1631 | Test acc: 95.00%\n",
      "Epoch 08 | Train loss: 0.1395 | Test loss: 0.1593 | Test acc: 95.02%\n",
      "Epoch 09 | Train loss: 0.1300 | Test loss: 0.1550 | Test acc: 95.40%\n",
      "Epoch 10 | Train loss: 0.1191 | Test loss: 0.1472 | Test acc: 95.50%\n"
     ]
    }
   ],
   "source": [
    "# Use GPU if available, otherwise use CPU\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "if device == \"cuda\": \n",
    "    print(\"Using GPU\") \n",
    "else:\n",
    "    print(\"Using CPU\")\n",
    "    \n",
    "# Instantiate the model\n",
    "model = NN_model().to(device)\n",
    "\n",
    "# Define the loss function and the optimizer\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-3)\n",
    "\n",
    "# Define the hyperparameters\n",
    "batch_size = 128\n",
    "num_epochs = 10\n",
    "\n",
    "# Create the training batches from the black and white training dataset\n",
    "train_loader = DataLoader(\n",
    "    mnist_bw,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    drop_last=True\n",
    "    # If the last batch is smaller than the batch size, we drop it\n",
    ")\n",
    "\n",
    "# Create the testing batches from the black and white testing dataset\n",
    "test_loader = DataLoader(\n",
    "    mnist_bw_test,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=False\n",
    ")\n",
    "\n",
    "# We could load the weights if we already trained the model\n",
    "# weights_path = \"nn_weights.pth\"\n",
    "# model.load_state_dict(torch.load(weights_path, map_location=device))\n",
    "\n",
    "\n",
    "# The main training loop\n",
    "# An epoch is one pass through the entire training dataset\n",
    "for epoch in range(num_epochs):\n",
    "    # We put the model in training mode\n",
    "    model.train()\n",
    "    train_loss = 0\n",
    "\n",
    "    # We iterate through all the batches (x are the images of a batch and y are the labels)\n",
    "    for x, y in train_loader:\n",
    "        x, y = x.to(device), y.to(device)\n",
    "\n",
    "        # We reset the gradient\n",
    "        optimizer.zero_grad()\n",
    "        # We predict the labels by doing a forward pass (it's a vector)\n",
    "        out = model(x)\n",
    "        # We comput the batch loss by comparing the predicted labels with the true labels\n",
    "        loss = criterion(out, y)\n",
    "        # We backpropagate the loss and update the weights\n",
    "        loss.backward()\n",
    "        # We take an optimization step to update the weights\n",
    "        optimizer.step()\n",
    "\n",
    "        # We accumulate the loss for all batches\n",
    "        train_loss += loss.item()\n",
    "\n",
    "    # We get the average batch loss\n",
    "    train_loss /= len(train_loader)\n",
    "\n",
    "    # We call the evaluate function to get the metrics\n",
    "    test_loss, test_acc = evaluate(\n",
    "        model, test_loader, criterion, device\n",
    "    )\n",
    "\n",
    "    print(\n",
    "        f\"Epoch {epoch+1:02d} | \"\n",
    "        f\"Train loss: {train_loss:.4f} | \"\n",
    "        f\"Test loss: {test_loss:.4f} | \"\n",
    "        f\"Test acc: {test_acc*100:.2f}%\"\n",
    "    )\n",
    "\n",
    "    torch.save(model.state_dict(), \"nn_weights.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e47d877a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Largest weight: 70 (must be <= 127)\n"
     ]
    }
   ],
   "source": [
    "# Quantization: multiply every weight and bias by 64 and round to an integer\n",
    "SCALE = 64\n",
    "\n",
    "# This line only takes the linear layers of the model (only those have weights)\n",
    "# And ingores all the other layers (like Flatten, ReLu, etc.)\n",
    "linears = [m for m in model.model if isinstance(m, nn.Linear)]\n",
    "\n",
    "# l.weight is the weight matrix, e.g. shape (32, 784) for layer 1\n",
    "# .detach() tells PyTorch you just want the numbers, not the gradient-tracking machinery\n",
    "# .cpu() moves them off the GPU if they were there\n",
    "# W are three matrices (we have three linear layers) of the weights\n",
    "W = [torch.round(l.weight.detach().cpu() * SCALE).to(torch.int32) for l in linears]\n",
    "# B are three vectors of the biases\n",
    "B = [torch.round(l.bias.detach().cpu() * SCALE).to(torch.int32) for l in linears]\n",
    "\n",
    "# Write the integer weights into a C header for the Arduino sketch\n",
    "with open(\"main.c/nn_weights.h\", \"w\") as f:\n",
    "    # Pragma once tells the compiler to include this file only once\n",
    "    # Even if instantiated in multiple files\n",
    "    f.write(\"#pragma once\\n\")\n",
    "    f.write(\"#include <stdint.h>\\n\\n\")\n",
    "\n",
    "    for layer in range(3):\n",
    "        weights = W[layer]\n",
    "        biases = B[layer]\n",
    "        name = layer + 1                 # C names start at 1: W1, W2, W3\n",
    "        num_neurons = weights.shape[0]   # rows\n",
    "        num_inputs = weights.shape[1]    # columns\n",
    "\n",
    "        # Weight matrix: one line of numbers per neuron\n",
    "        f.write(f\"const int8_t W{name}[{num_neurons}][{num_inputs}] = \")\n",
    "        f.write(\"{\\n\")\n",
    "        for row in weights:\n",
    "            numbers = [str(x) for x in row.tolist()]\n",
    "            f.write(\"  {\" + \", \".join(numbers) + \"},\\n\")\n",
    "        f.write(\"};\\n\")\n",
    "\n",
    "        # Bias list: all on one line\n",
    "        numbers = [str(x) for x in biases.tolist()]\n",
    "        f.write(f\"const int32_t BIAS{name}[{num_neurons}] = \")   # not B1: Arduino already defines B1, B2, ...\n",
    "        f.write(\"{\" + \", \".join(numbers) + \"};\\n\\n\")\n",
    "\n",
    "# Check that every weight fits in int8 (-128 to 127)\n",
    "largest = 0\n",
    "for weights in W:\n",
    "    largest = max(largest, weights.abs().max().item())\n",
    "print(\"Largest weight:\", largest, \"(must be <= 127)\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "urob",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
