2-mnist_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Neural Network Hands-On Tutorial Part 2\n",
    "\n",
    "## MNIST Classification Task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "# Import necessary libraries\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import DataLoader\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Set device (GPU if available, else CPU)\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Loading the MNIST Dataset\n",
    "\n",
    "`torchvision` provides many built-in datasets, which are common benchmarks and can be used for learning purposes and drafting your own neural network implementations.\n",
    "\n",
    "Run the next cell to download the MNIST dataset and store it to the `./data` folder.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True)\n",
    "test_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Looking at the Dataset\n",
    "\n",
    "The MNIST train dataset contains 60,000 pairs of data in the shape of (image, label).\n",
    "\n",
    "Each image is a grayscale image cropped to 28*28 pixels, with a centered handwritten digit.\n",
    "\n",
    "_Note_: the test dataset contains 10,000 images with the same shape as the train data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "60000\n",
      "(<PIL.Image.Image image mode=L size=28x28 at 0x107CC3400>, 5)\n"
     ]
    }
   ],
   "source": [
    "print(len(train_dataset))\n",
    "print(train_dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Looking at the Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAjIAAACHCAYAAAAMVLO2AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/H5lhTAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAYz0lEQVR4nO3dfVjV1QEH8O9V4n0qFr6QCeR7NkJJU3OiUW7TxFyluWVqTVGz1BVYi5mWOMTUmWawnLqQTdlSLPf00JY4X3LO1zZNkkxooI/vQIBC6tkfPp7O+SFwL17u5Vy+n+fxeb6/+zv3dw8cwePvvPxsQggBIiIiIgM1c3cFiIiIiOqLHRkiIiIyFjsyREREZCx2ZIiIiMhY7MgQERGRsdiRISIiImOxI0NERETGYkeGiIiIjMWODBERERnLpR2ZtWvXwmazIT8/3+H3Dh48GPfee69T6xMWFoYJEyY49ZqejO1nNraf+diGZmP7NQzekamn/Px82Gy2m/5Zv369u6tHdrh27RpSUlIQHh4OX19fRERE4M9//rO7q0X1kJGRAZvNhsDAQHdXhRyQlJSE2NhYtG3bFjabDXPnznV3lcgBX331FZ544gkEBQXB398fAwcORE5Ojsvr4eXyT/QwY8eOxbBhw7TX+vfv76bakCNee+01JCcnY9KkSejTpw82b96Mn//857DZbHjqqafcXT2yU1lZGRISEhAQEODuqpCDEhMT0a5dO/Tq1QvZ2dnurg454H//+x/69++P5s2bIz4+HgEBAVizZg2GDh2KTz/9FIMGDXJZXdiRuUW9e/fG008/7e5qkIOKioqwePFiPP/881ixYgUA4Je//CWio6MRHx+PJ598Es2bN3dzLcke8+fPxw9+8AMMGTIEWVlZ7q4OOeDEiRMICwvDuXPnEBwc7O7qkAOSk5NRXFyMw4cPo1u3bgCASZMmoXv37pg1axb279/vsrq4fWhp8+bNGD58OEJCQuDj44NOnTrhzTffxNWrV29afv/+/RgwYAD8/PwQHh6O1NTUamUqKyvx+uuvo3PnzvDx8cFdd92FhIQEVFZW1lmf48eP4/jx4w59DeXl5aiqqnLoPZ7C1PbbvHkzvvvuO0ybNk2+ZrPZMHXqVBQWFmL37t11XsMTmNp+N+Tl5WHp0qVYsmQJvLya5v/LTG7DsLAwu8p5MlPbb8eOHejVq5fsxACAv78/YmNjceDAAeTl5dV5DWdx+0/+2rVrERgYiF/96lcIDAzE1q1bMWfOHJSWlmLRokVa2YsXL2LYsGEYPXo0xo4di8zMTEydOhXe3t549tlnAVyf9xAbG4udO3di8uTJ6NGjB/773/9i6dKlOHbsWJ3/Y4uJiQEAuydjzZs3D/Hx8bDZbIiKikJSUhKGDh3q8PfBVKa238GDBxEQEIAePXpor/ft21eeHzhwoAPfCTOZ2n43zJw5E0OGDMGwYcOQmZnp8NfvCUxvw6bO1ParrKxEUFBQtdf9/f0BXO9wdenSxc7vwi0SLrRmzRoBQJw4cUK+VlFRUa1cXFyc8Pf3F5cvX5avRUdHCwBi8eLF8rXKykoRGRkp2rRpI6qqqoQQQqSnp4tmzZqJHTt2aNdMTU0VAMSuXbvka6GhoWL8+PFaudDQUBEaGlrn11JQUCCGDh0q3n33XfHhhx+K3/3ud6Jjx46iWbNmYsuWLXW+30Se1H7Dhw8Xd999d7XXy8vLBQDxyiuv1HkN03hS+wkhxJYtW4SXl5c4cuSIEEKI8ePHi4CAALveaypPa8Mbzp49KwCI119/3aH3mcaT2m/EiBGiVatWorS0VHu9f//+AoB466236ryGs7h9aMnPz0/mb7/9FufOncOPfvQjVFRUIDc3Vyvr5eWFuLg4eezt7Y24uDicOXNGjsf95S9/QY8ePdC9e3ecO3dO/nnooYcAoM4Z1fn5+Xb9T6Jjx47Izs7GlClTMGLECMyYMQMHDx5EcHAwXnrpJXu/fOOZ2n6XLl2Cj49Ptdd9fX3l+abA1ParqqrCrFmzMGXKFNxzzz32frkeydQ2pOtMbb+pU6eiuLgYY8aMwcGDB3Hs2DHMnDkT+/btA+Da36FuH1o6cuQIEhMTsXXrVpSWlmrnSkpKtOOQkJBqKxO6du0K4Po3v1+/fsjLy8PRo0drnDh25swZJ9Ze17p1a0ycOBHJyckoLCxEhw4dGuyzGgtT28/Pz++m48WXL1+W55sCU9tv6dKlOHfuHObNm+eU65nM1Dak60xtv5/+9KdYvnw5XnnlFfTu3RsA0LlzZyQlJSEhIcGlWyG4tSNTXFyM6OhotGjRAm+88QY6deoEX19fHDhwALNnz8a1a9ccvua1a9fwwx/+EEuWLLnp+bvuuutWq12rG9e/cOGCx3dkTG6/9u3bIycnB0II2Gw2+fqpU6cAXP+F4elMbb+SkhLMnz8f06ZNQ2lpqfzlX1ZWBiEE8vPz4e/vjzZt2tzyZzV2prYhXWd6+02fPh0TJ07Ef/7zH3h7eyMyMhJ/+MMfAHzfwXIFt3Zktm3bhvPnz2Pjxo3amvMTJ07ctPzJkydRXl6u9UiPHTsG4PvZ7506dcLnn3+OmJgY7R8oV/n6668BoEksJTS5/SIjI7Fq1SocPXpUG5rYs2ePPO/pTG2/ixcvoqysDCkpKUhJSal2Pjw8HCNHjmwSS7FNbUO6zhPaLyAgQNs77R//+Af8/Pzw4IMPNvhn3+DWOTI39ukQQsjXqqqqsHLlypuWv3LlCtLS0rSyaWlpCA4ORlRUFABg9OjRKCoqwnvvvVft/ZcuXUJ5eXmtdbJ36dnZs2ervVZUVITVq1cjIiIC7du3r/MapjO5/UaOHInbbrtNq6sQAqmpqbjzzjsxYMCAOq9hOlPbr02bNti0aVO1P0OGDIGvry82bdqEV199tdZreApT25Cu87T2++yzz7Bx40Y899xzaNmyZb2uUR9uvSMzYMAABAUFYfz48XjxxRdhs9mQnp6uNaoqJCQECxcuRH5+Prp27YoNGzbg0KFD+P3vf4/bbrsNADBu3DhkZmZiypQpyMnJwYMPPoirV68iNzcXmZmZyM7Oxv33319jnexdepaQkIDjx48jJiYGISEhyM/PR1paGsrLy7Fs2bL6fUMMY3L7dejQATNnzsSiRYvw3XffoU+fPsjKysKOHTuQkZHRJDbDM7X9/P398dhjj1V7PSsrC//+979ves5TmdqGN6Snp6OgoAAVFRUAgO3bt2P+/PmyHqGhoY58O4xjcvsVFBRg9OjRiI2NRbt27XDkyBGkpqYiIiICCxYsqN83pL5ctj5K3Hzp2a5du0S/fv2En5+fCAkJEQkJCSI7O1sAEDk5ObJcdHS06Nmzp9i3b5/o37+/8PX1FaGhoWLFihXVPqeqqkosXLhQ9OzZU/j4+IigoCARFRUl5s2bJ0pKSmS5W1l69qc//UkMGjRIBAcHCy8vL3HHHXeIUaNGif379zv6bTGGJ7WfEEJcvXpVLFiwQISGhgpvb2/Rs2dPsW7dOke+JUbxtPazaqrLr01uwxtLim/2R627p/Ck9rtw4YIYOXKkaNeunfD29hbh4eFi9uzZ1ZZju4JNiBq6fkRERESNnNv3kSEiIiKqL3ZkiIiIyFjsyBAREZGx2JEhIiIiY7EjQ0RERMZiR4aIiIiMZfeGeNyq2j2ctTqe7ecebD+zOXN3Crahe/Bn0Gz2tB/vyBAREZGx2JEhIiIiY7EjQ0RERMZiR4aIiIiMxY4MERERGYsdGSIiIjIWOzJERERkLHZkiIiIyFjsyBAREZGx2JEhIiIiY7EjQ0RERMZiR4aIiIiMZfdDI4kag6ioKO14+vTpMj/zzDMyv//++1q55cuXy3zgwIEGqh0REbka78gQERGRsdiRISIiImOxI0NERETGsgkhhF0FbbaGrsstad68uXbcsmVLu96nzrHw9/eXuVu3blq5559/Xua33npL5rFjx2rlLl++LHNycrJ2bt68eXbVSWVn89SpsbdfbSIjI2XeunWrdq5FixZ2XaOkpETm22+/3Sn1sgfbz/liYmJkzsjI0M5FR0fL/OWXX97yZzmr/YCm14aJiYnasfr7r1mz7/8PPXjwYK3cP//5T6fWgz+DZrOn/XhHhoiIiIzFjgwREREZq1Euv+7YsaPM3t7e2rkBAwbIPHDgQJlbtWqllXv88cdvqQ6FhYXa8dtvvy3zqFGjZP7222+1cp9//rnMzr5F2pT07dtX5g8++EBm65ChettRbYuqqiqtnDqc1K9fP5mtS7Gt7zPZoEGDtGP1e7Bp0yZXV8dp+vTpI/PevXvdWBOymjBhgsyzZ8/Wzl27du2m73Hm8B01TbwjQ0RERMZiR4aIiIiM1SiGltRVKYC+MsXe1UfOoN76tM64Lysrk1ldKXHq1Cmt3MWLF2V2xqoJT6auEuvdu7d2bt26dTK3b9/eruvl5eXJnJKSop1bv369zLt27ZLZ2s6//e1v7fosE1hXg3Tp0kVm04aW1FUu4eHhMoeGhmrluLLEvdT28PX1dWNNmp4HHnhAO3766adlVlfz9ezZs8ZrvPzyyzKfPHlSO6dO5VB/P+/Zs8fxyjoZ78gQERGRsdiRISIiImOxI0NERETGahRzZL755hvt+Pz58zI7Y46MOoZXXFysnRsyZIjM6tLb9PT0W/5cql1aWprM1h2S60OdZxMYGKidU5fCq3NHIiIibvlzGyv1aeAAsHv3bjfV5Nap86QmTZokszpWDwC5ubkuqxNd9/DDD8v8wgsv1FhObZtHH31U5tOnTzdMxZqAMWPGyLxs2TLt3B133CGzOnds27ZtWrng4GCZFy1aVONnqddQ3/PUU0/ZX+EGwjsyREREZCx2ZIiIiMhYjWJo6cKFC9pxfHy8zOotSAA4ePCgzOpuu1aHDh2S+ZFHHpG5vLxcK6cuRZsxY4Z9FaZ6i4qKknn48OEy17ZsVh0W+uijj7Rz6gM81eWC6t8TQF8W/9BDD9n1uaZTlyybbtWqVTd9XV1yT66hLsMFgDVr1shc21QAddiioKDA+RXzUF5e+j/T999/v8zvvfeezOp2FgCwfft2md98802Zd+7cqZXz8fGROTMzU+ahQ4fWWKd9+/bVVW2X8pzfdERERNTksCNDRERExmJHhoiIiIzVKObIWGVlZcmsPq4A0J9wfN9998n83HPPaeXUuRPWeTGqI0eOyDx58mSH60q1sz5+4u9//7vMLVq0kNn6BNyPP/5YZnVptrrVNqA/YkCdR3H27FmtnPpUcvVRFOo8HUBfwm19MrYJ1OXkbdu2dWNNnKumuRfq3ydyjfHjx2vHISEhNy1nXeb7/vvvN1SVPJr6qAGg5vli1p8FdWl2aWlpjddXy9U2L6awsFDmP/7xjzWWcwfekSEiIiJjsSNDRERExmqUQ0uq2m6JlZSU1HhO3f1zw4YNMqvDCtQwunbtKrO6lB7QhwjOnTsns/Up4uqtS/XJ43/729+0ctZjR/n5+WnHL730ksy/+MUvbuna7jBs2DCZrV+bSazDYuoTr1VFRUWuqE6Tp+4S++yzz2rn1N+p6s7p8+fPb/B6eSp1ufSvf/1r7Zw6DL9y5UqZ1WF2oPZ/O1WvvfaaXeVefPFFma1D9+7GOzJERERkLHZkiIiIyFiNfmipNnPnzpVZ3TEW0Fe3qA81++STTxq8Xk2NujMkoK8YU4c6AH3VmfpQQ+tOke4aFunYsaNbPtdZunXrVuM5dYVeY6f+HQL0oaZjx47JrP59IucKCwuT+YMPPrDrPcuXL5c5JyfH2VXyWHPmzNGO1eEk9WHGAJCdnS3z7NmzZb506VKN1/f19ZXZujJJ/Z2n7nRuHRrcvHlzjdd3N96RISIiImOxI0NERETGYkeGiIiIjGX0HBl1x151uTWg78qqPiHUOm6rzs145513ZLbuNEs169Wrl3ZsnRejGjlypMzqU62p4e3du9fdVdB2cwaAn/zkJzKrO5jWtsOoujRVXe5LzqW2jbpjtNWnn34q87Jlyxq0Tp6kVatWMk+bNk07p/77o86JAYDHHnvMrut37txZ5oyMDJmt80lVf/3rX2VOSUmx63MaA96RISIiImOxI0NERETGMnpoSXX8+HHteMKECTKvWbNG5nHjxmnl1OOAgACZrQ84s+48S99bsmSJdqwu4bMOHzWG4aRmzb7vvzelnZ5bt27t8HvUB7MCetuq2xp06NBBK+ft7S2zukOy+r0H9CWje/bskbmyslIr5+X1/a+q/fv321V3cox1yCI5Ofmm5Xbu3Kkdqw+RrG23ddKpPyPqzslW6o66ANCmTRuZJ06cKHNsbKxW7t5775U5MDBQZuu0CfV43bp1Mtf2sOXGhndkiIiIyFjsyBAREZGxPGZoyWrTpk0y5+XlyWwdBomJiZF5wYIFMoeGhmrlkpKSZOaD6oBHH31U5sjISO2ceqvyww8/dFWV7KYOJ1lvsx46dMjFtXEudajG+rWlpqbKbH0QXU2sq1XUoaUrV67IXFFRoZX74osvZF69erXM1h2c1aHG06dPy1xYWKiVU3d6zs3NtavuVLf67N779ddfa8dqu5H91B17rQ9hDA4OlvnEiRPaOXtX1J48eVJm9QGS7du318qpD+/96KOP7Lp2Y8M7MkRERGQsdmSIiIjIWOzIEBERkbE8do6M6vDhwzKPHj1aOzdixAiZ1WXacXFxWrkuXbrI/Mgjjzi7isZR5yyoywgB4MyZMzJv2LDBZXVSWZ/IrT4pXbV161bt+NVXX22oKrmEukNoQUGBdm7AgAEOX++bb77RjrOysmQ+evSozP/6178cvrbV5MmTZVbnCADV52WQc6hPT7Z3K4KalmWTY9Rdqa1L37ds2SKzddsEdasR9YnUa9eu1cpduHBB5vXr18tsnSOjnjMV78gQERGRsdiRISIiImM1iaEllfUhc+np6TKvWrVKZnUnUQAYNGiQzIMHD5Z527ZtTq2fJ1B3ZXXljsjqcFJiYqJ2Lj4+XmZ1ae/ixYu1cmVlZQ1UO9dbuHChu6vgEHUrBCt7lwZT3dTtEmp7OKdKHcL48ssvnV2lJk/d1RqoPrRaH+q/WdHR0TJbhxA9YdiWd2SIiIjIWOzIEBERkbGaxNCSujvpE088oZ3r06ePzNbhJJW6U+n27dudWDvP48rdfNXb5Orw0ZgxY7Ry6q3xxx9/vMHrRc6l7tRNt+aTTz6ROSgoqMZy6ko09SG8ZAZ1ZWltu5lz1RIRERGRG7EjQ0RERMZiR4aIiIiM5TFzZLp166YdT58+Xeaf/exnMrdr186u6129elU7VpcR27sDpidTn4KsZkDfpXLGjBlO/dxZs2Zpx7/5zW9kbtmypcwZGRlauWeeecap9SAy1e233y5zbb/LVq5cKbMnbUvQVGRnZ7u7Ci7DOzJERERkLHZkiIiIyFjGDS2pQ0Njx46VWR1KAoCwsDCHr71v3z6Zk5KStHOuXFJsAnUJn3U5n9pGb7/9tnZu9erVMp8/f17mfv36aeXGjRsn83333Sdzhw4dtHLqQw3VW6nqbXEyj3W4smvXrjI74wGVTYn6MFwAaNbMvv+/fvbZZw1RHXKRH//4x+6ugsvwjgwREREZix0ZIiIiMlajHFpq27atzPfcc492bsWKFTJ3797d4WtbH861aNEimdXdX7kyqf6aN28u87Rp07Rz6q66paWlMnfp0sWua1tvd+fk5Mg8Z84ch+pJjZd1uNLe4RC6Tt3x+uGHH9bOqb/bqqqqZH7nnXe0cqdPn26YypFL3H333e6ugsvwtwMREREZix0ZIiIiMhY7MkRERGQst82Rad26tcxpaWnaOXV8t77jfOpcisWLF8ts3e3w0qVL9bp+U7d7926Z9+7dq51TnyhupS7NVudCWalLs9Wnszp7p2AyQ//+/WVeu3at+ypiiFatWslc227mRUVFMr/88ssNWSVysR07dsiszjHzxPmfvCNDRERExmJHhoiIiIzVoENLDzzwgHYcHx8vc9++fWW+884763X9iooKma07yC5YsEDm8vLyel2falZYWCiz+lBOAIiLi5M5MTHRrustW7ZMO3733Xdl/uqrr+pTRTKYdWdfInLM4cOHZc7Ly5PZOl2jU6dOMp89e7bhK9YAeEeGiIiIjMWODBERERmLHRkiIiIyVoPOkRk1alStxzX54osvZN6yZYt27sqVKzKry6qLi4vrUUNyhlOnTmnHc+fOvWkmqs3HH38s85NPPunGmpgvNzdXZutjPQYOHOjq6pCbqXNGV61apZ1LSkqS+YUXXpBZ/Xe4seMdGSIiIjIWOzJERERkLJuwPma2poJcDukWdjZPndh+7sH2M5uz2g9gG7oLfwaBFi1ayJyZmamdU5+OvnHjRpknTpyolXPXNib2tB/vyBAREZGx2JEhIiIiY3FoqZHjbVGzsf3MxqEl8/FnUKcOMwH6qqWpU6fKHBERoZVz1yomDi0RERGRR2NHhoiIiIzFjgwREREZi3NkGjmO75qN7Wc2zpExH38GzcY5MkREROTR2JEhIiIiY9k9tERERETU2PCODBERERmLHRkiIiIyFjsyREREZCx2ZIiIiMhY7MgQERGRsdiRISIiImOxI0NERETGYkeGiIiIjMWODBERERnr/zsv7C6/k9cuAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 700x250 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "idx = 0  # Change the index to see different images\n",
    "\n",
    "# Show some images\n",
    "fig, axes = plt.subplots(1, 5, figsize=(7, 2.5))\n",
    "for i in range(5):\n",
    "    axes[i].imshow(train_dataset[i+idx][0], cmap='gray')\n",
    "    axes[i].set_title(f\"label: {train_dataset[i+idx][1]}\")\n",
    "    axes[i].axis('off')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Looking at the Dataset\n",
    "\n",
    "Currently, the images are `PIL` images and the amplitudes range from 0 to 255."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "255 0 (28, 28)\n"
     ]
    }
   ],
   "source": [
    "img = train_dataset[0][0]\n",
    "print(np.max(img), np.min(img), np.shape(img))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Data Preprocessing\n",
    "\n",
    "In order to train a neural network with the dataset, the train dataset needs to be pre-processed. In `PyTorch` and `torchvision`, this can be achieved by `torchvision.transforms`, which includes many common image processing methods.\n",
    "\n",
    "Note that the images in the MNIST dataset are already _centered_ and _cropped to the same shape_. (For your own dataset, remember to perform these steps.)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "For the MNIST dataset, we only need to perform two steps:\n",
    "\n",
    "1. Convert the PIL images with amplitude $[0,255]$ to PyTorch Tensors in $[0,1]$, with `transforms.ToTensor()`\n",
    "2. Normalize the images to $\\mu=0.5$ and $\\sigma=0.5$, with `transforms.Normalize()`\n",
    "\n",
    "Multiple transformations can be chained by using `transforms.Compose`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist_transform = transforms.Compose([\n",
    "    transforms.ToTensor(),  # Convert image to tensor\n",
    "    transforms.Normalize((0.5,), (0.5,))  # Normalize image to mean 0.5 and std 0.5\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Let's define the train and test dataset again, with the proper transformations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=mnist_transform)\n",
    "test_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=mnist_transform)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "We can load mini-batches of data from the dataset using `DataLoader`.\n",
    "\n",
    "Setting `shuffle=True` allows the batches to be sampled in a random order across different episodes. \n",
    "\n",
    "Why do we need it? $\\rightarrow$ Prevents converging into local optima and overfitting.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "batchsize = 64\n",
    "train_loader = DataLoader(dataset=train_dataset, batch_size=batchsize, shuffle=True)\n",
    "test_loader = DataLoader(dataset=test_dataset, batch_size=batchsize, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "We can now look at one batch of data sampled from the `DataLoader`\n",
    "\n",
    "The input data shape would be: [batch_size, channel, height, width]. Here channel=1 because we are using grayscale images. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input batch shape:  torch.Size([64, 1, 28, 28])\n",
      "Output batch shape:  torch.Size([64])\n"
     ]
    }
   ],
   "source": [
    "images, labels = next(iter(train_loader))\n",
    "print(\"Input batch shape: \", images.shape)\n",
    "print(\"Output batch shape: \", labels.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take a look at the sampled dataset (run the next cell several times to see random samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAjIAAACHCAYAAAAMVLO2AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/H5lhTAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAZoUlEQVR4nO3de1RVVR4H8C/KQwgmFfGBA+jIKI6PLNQRNbFAcwxUzErIfDY+ysfMOOFq1LR8pDNjZg8SscEW5FLKxrEcBROM1LGU0p5IkfhAW8qYL1BROfOHyz17H7nXC9zLvfvy/azlWr/D2ffczd0e2Jy99297GIZhgIiIiEhDjZxdASIiIqLaYkeGiIiItMWODBEREWmLHRkiIiLSFjsyREREpC12ZIiIiEhb7MgQERGRttiRISIiIm2xI0NERETaqteOzLp16+Dh4YGSkpIav3bgwIHo2rWrXevTrl07jB8/3q7XdGdsP72x/fTHNtQb288x+ESmlk6ePIkxY8agU6dOCAgIQNOmTdG7d2+8/fbb4K4Prq+kpAQeHh7V/tuwYYOzq0c2Ki4uRlJSElq2bAlfX1/8+te/xty5c51dLbJBYWEhkpOT0aNHDwQEBKBNmzZ4+OGHceDAAWdXjWy0ZMkSDBs2DK1atYKHhwcWLlzolHp4OuVd3UBZWRlOnDiBUaNGITQ0FNeuXcOOHTswfvx4HD58GEuXLnV2FckGiYmJGDp0qPK1qKgoJ9WGauLgwYMYOHAg2rZti9mzZyMwMBDHjh3D8ePHnV01ssHatWvx1ltv4ZFHHsHTTz+N8+fPIzU1FX369MH27dsRGxvr7CrSHcybNw+tW7fGvffei+zsbKfVgx2ZWurevTt27dqlfG369OmIj4/Hq6++ikWLFqFx48bOqRzZ7L777sOYMWOcXQ2qoaqqKjz55JOIiIhAXl4efH19nV0lqqHExEQsXLgQ/v7+4msTJ05E586dsXDhQnZkNHDkyBG0a9cOZWVlCAoKclo9nD609K9//QsPP/wwgoOD4ePjgw4dOmDRokW4ceNGteULCgrQt29f+Pr6on379li9evVtZa5evYoFCxYgPDwcPj4+CAkJQXJyMq5evXrH+hQXF6O4uLjW30+7du1QUVGBysrKWl9DJ+7QfuXl5Q2mvcx0bb+cnBx8/fXXWLBgAXx9fVFRUWGxzu5O1zaMjIxUOjEAEBgYiPvvvx/ffffdHV/vLnRtP+Dm7ztX4PQnMuvWrYO/vz/+9Kc/wd/fH7m5uXj++edx4cIF/O1vf1PK/vzzzxg6dCgee+wxJCYmIisrC9OmTYO3tzcmTpwI4OZfasOGDcPu3bsxefJkdO7cGV999RVWrlyJoqIibN682Wp9YmJiAMDmyViXL19GeXk5Ll26hI8//hjp6emIiopqMH8h6t5+L7zwAp599ll4eHggMjISS5YsweDBg2v8OehK1/b76KOPAAA+Pj7o2bMnCgoK4O3tjYSEBKSkpKB58+a1+0A0pGsbWvLTTz+hRYsWtXqtjtyt/ZzCqEfp6ekGAOPIkSPiaxUVFbeVmzJliuHn52dcuXJFfC06OtoAYKxYsUJ87erVq0aPHj2Mli1bGpWVlYZhGEZGRobRqFEj45NPPlGuuXr1agOAsWfPHvG1sLAwY9y4cUq5sLAwIywszObv6aWXXjIAiH8xMTHGsWPHbH69Ttyp/Y4ePWoMHjzYePPNN40tW7YYr7zyihEaGmo0atTI+PDDD+/4eh25U/sNGzbMAGAEBgYaTzzxhPHee+8Z8+fPNzw9PY2+ffsaVVVVd7yGjtypDauTn59veHh4GPPnz6/V612du7bfmTNnDADGggULavQ6e3H60JL85OLixYsoKyvD/fffj4qKChQWFiplPT09MWXKFHHs7e2NKVOm4PTp0ygoKAAAvPvuu+jcuTMiIiJQVlYm/j344IMAgLy8PKv1KSkpqVFPNDExETt27MD69euRlJQE4OZTmoZC1/YLDQ1FdnY2pk6divj4eMyaNQtffPEFgoKCMHv2bFu/fe3p2n6XLl0CAPTq1QuZmZl45JFH8OKLL2LRokXYu3cvdu7cadP37w50bUOz06dPIykpCe3bt0dycnKNX68rd2k/Z3L60NI333yDefPmITc3FxcuXFDOnT9/XjkODg7GXXfdpXytY8eOAG5++H369MH333+P7777zuLEo9OnT9ux9kBYWBjCwsIA3OzUTJ48GbGxsTh8+HCDGF7Svf1kzZs3x4QJE7Bs2TKcOHECv/zlLx32Xq5C1/a7dW8lJiYqX09KSsJzzz2HvXv3NpjJorq2oay8vBxxcXG4ePEidu/efdvcGXfmDu3nbE7tyJw7dw7R0dH4xS9+gRdffBEdOnRAkyZN8Pnnn2POnDmoqqqq8TWrqqrQrVs3vPzyy9WeDwkJqWu1rRo1ahTS0tKQn5+Phx56yKHv5Wzu2H63rn/27Fm378jo3H7BwcEAgFatWilfb9myJYCbcwkaAp3b8JbKykqMHDkSX375JbKzs+2e9M2VuUP7uQKndmR27dqF//73v3j//fcxYMAA8fUjR45UW/7kyZMoLy9XeqRFRUUA/j97ukOHDjh06BBiYmLg4eHhuMpbcGtYydyTdkfu2H4//vgjADh1KWF90bn9IiMjkZaWhtLS0tvqCDSM9gP0bkPg5i/dsWPHYufOncjKykJ0dLRD38/V6N5+rsKpc2Ru5VkxpEy4lZWVSElJqbb89evXkZqaqpRNTU1FUFAQIiMjAQCPPfYYSktLkZaWdtvrb60wssbWpWdnzpyp9utvvfUWPDw8cN99993xGrpzt/YrLS3FP/7xD3Tv3h1t2rS54zV0p3P7DR8+HD4+PkhPT1f+al27di0AYNCgQXe8hjvQuQ0BYMaMGdi4cSNSUlIwcuRIm17jTnRvP1fh1Ccyffv2RbNmzTBu3DjMnDkTHh4eyMjIsJjiPzg4GMuXL0dJSQk6duyIjRs34uDBg1izZg28vLwAAE8++SSysrIwdepU5OXloV+/frhx4wYKCwuRlZWF7Oxs9OzZ02KdbF16tmTJEuzZswdDhgxBaGgozp49i02bNmH//v2YMWMGwsPDa/ehaETn9ktOTkZxcTFiYmIQHByMkpISpKamory8HKtWrardB6IZnduvdevWmDt3Lp5//nkMGTIEI0aMwKFDh5CWlobExET06tWrdh+KZnRuw1deeQUpKSmIioqCn58fMjMzlfMJCQm3zQdxNzq3HwBkZGTg6NGjqKioAADk5+dj8eLFoh635o86XH0ukapu6dmePXuMPn36GL6+vkZwcLCRnJxsZGdnGwCMvLw8US46Otro0qWLceDAASMqKspo0qSJERYWZrz++uu3vU9lZaWxfPlyo0uXLoaPj4/RrFkzIzIy0njhhReM8+fPi3J1WXqWk5NjxMXFGcHBwYaXl5cREBBg9OvXz0hPT29QSz91bb/169cbAwYMMIKCggxPT0+jRYsWRkJCglFQUFDTj0Ub7tR+hmEYVVVVxmuvvWZ07NjR8PLyMkJCQox58+aJZajuyJ3acNy4cUrqCvM/+Xt0F+7UfrfqZKn95Lo7modhcIdDIiIi0pPT88gQERER1RY7MkRERKQtdmSIiIhIW+zIEBERkbbYkSEiIiJtsSNDRERE2rI5IV5DSXXsauy1Op7t5xxsP73ZMzsF29A5eA/qzZb24xMZIiIi0hY7MkRERKQtdmSIiIhIW+zIEBERkbbYkSEiIiJtsSNDRERE2rJ5+TURkaxJkyYiDggIUM5t3bpVxD179hSxeQnrtm3bRBwfHy/iGzdu2K2eROTe+ESGiIiItMWODBEREWnLw7Ax7SGzGjoHs1Lqzd3aTx5C2rhxo4gfeughi685d+6ciJs2bWqxXL9+/US8b9++2lXQzpjZV3/udg82NMzsS0RERG6NHRkiIiLSFlctEZHNnnrqKRHLw0nXr19Xyo0ePVrEZ8+eFfH777+vlJOHmiZPnixiVxlackc9evQQsbxSzKyoqEjE8jAiNUxDhgwR8YYNG0QcHh6ulCsrK6u3Ot3CJzJERESkLXZkiIiISFvsyBAREZG2XH6OTEREhHL8wAMPiHjw4MEiPnXqlFKuTZs21V6vf//+ynFgYKCIv/76axFHRUUp5crLy22sMbkCeSkvAIwYMULEQUFBDn3v8ePHO/T69alFixbK8TPPPFNtufXr1yvH//znP6stN2bMGOU4MzNTxCUlJbWoId1JRkaGcizPX2rcuLHF133zzTcirqqqEvG7775rx9qRq4qNjVWO5XlSchqGqVOnKuUWL17s2IpVg09kiIiISFvsyBAREZG2XD6zr/kR9fDhw+vlfYcNG6Ycf/jhh/XyvmbMSgn4+vqKOCQkRDknDy/OmTNHxK1bt1bKWXqEfvr0aeVYzkJr9s4774h406ZNIr5w4YJS7sSJEyLWvf3kZZYA8Oijj1ZbrkOHDsqxrcNE8tJN+XO7cuWKjTV0LF0z+8rDSU888USdrydv4vnqq68q5z799FMRZ2Vl1fm97E33e7A+yRvBmn/3yukW5Ptz5syZSrm1a9fatU7M7EtERERujR0ZIiIi0pbLr1oyq6ysFPG3334r4pycHKXc5cuXRSw/EjQ/ppoyZYqI5eEIeTUTOZ6Xl5dyLLeLnPG1a9euFq9x7do1EZuHJrZt2yZieVgoPz9fKWde/dYQyZlfzUOssq+++krEtf3cfvjhh1q9jlTmx/vWhpMKCwtFnJeXZ7Hc448/LuLmzZuL+I9//KNS7scffxSxnJH52LFjVmpMrkhecWltI1h5Fa+9h5Jqg09kiIiISFvsyBAREZG22JEhIiIibbn88mt5x01Azb4rL9e0lby8DAC+/PJLEctLQdu2bauUc9bciYaydNA87r5ixYpqy8lznwBgzZo1Il61apWIXSVLrA7tJ89/AICdO3eKuHv37so5Odvrc889J+KtW7c6qHbO5crLr+W5TLm5uco5eVdxeU4MoM59OH78uMXryxmwly1bJuIJEyZYfI08b8q8s7az5szocA86U6dOnUQszyVs166dUk7+2Ttx4kQRO3pndC6/JiIiIrfGjgwRERFpy+WXX2/fvr3O15CX9q5cuVI5Jw8nlZWVidicrZXsb8CAASI2Dy3J5GGigQMHKue4xLPuzNmSzcNJskOHDonYXYeTdCEP3chDSdbKAdaHk2RnzpwRsZzN13w9eWPRbt26iVge+gJ4r7oK83SNN954Q8Tm4STZ/v37Rezo4aSa4hMZIiIi0hY7MkRERKQtdmSIiIhIWy4/R6a2/Pz8RCwvy500aZLF16SmpopYTsFMtTdmzBjlePbs2SLu0qWLiOW5FwDw2WefiXjPnj0i5ji7/T3zzDM2l7XHnDWyj1mzZon48OHDyjk5rYR5h/fakO/P77//Xjknz5GRyfc3AGzZsqXO9aC6Gz16tHLcvn37asvJP4MBYPDgwQ6rU13xiQwRERFpix0ZIiIi0pbLZ/a11e9+9zvleMGCBSLu3bu3Tdf41a9+JWJmhrVdq1atlGN5qWZsbKxyrlmzZiI+ePCgiGfMmKGUk4eTdOaq7SenJHjvvfeUc3FxcSL+4IMPlHNjx44VcUNIUeDKmX3lneDl1BEA8NNPP9n1vWSDBg1SjtevXy/iwMBAEcvDW8Dty7Hri6veg/VJTm/x17/+VTnXuHFjEZeWloo4JiZGKVdUVOSg2lnHzL5ERETk1tiRISIiIm1pt2opLCxMxPJGZgkJCUo5b2/vGl973bp1IjZnP7xy5UqNr9dQHD16VDmWP3t5+AgA/vKXv4g4MzNTxFwlVr+mTZsmYnkoyUxe8Qc0jOEkXcgb6NYn8yolHx+fasuZVzdR/WrTpo2I5ZWJ8lCS2YYNG0TsrKGk2uATGSIiItIWOzJERESkLXZkiIiISFvazZFJSUkRsXnJdV3JuzE3adJEOcc5MpaZ5yPJ8yiefvpp5dy+ffvqpU5kXXBwsLOrAEC9hx988EERDx061OJrduzYIeLc3FzlHLPHOt706dOVY39//2rLZWRk1Ed1yILFixeLWE4tYrZr165qX6MTPpEhIiIibbEjQ0RERNrSbmjp/Pnz1cZyNllAzQwrb35lXjq4f/9+Ed99990ivvfee5VyeXl5tayx+8vJyVGO5c3FPvnkE+XcO++8I2I5w+S3337roNqRMwwcOFDEc+fOFbGcjRZQMz3L2YatiYiIELG8jBwAVq9eLeKXXnpJOefIbLe6atmypYg7duyonJOzBRcWFtp0vUuXLon47Nmzdawd1UT//v2V41GjRolYzkr8888/K+XklBjy71Sd8IkMERERaYsdGSIiItKWdkNLSUlJdXr9uXPnlGN5c8h77rlHxHIGYbLOnAVZXjkiZ5cE1FUq8fHxIjZvXPjaa6+J2FkZTMn6cI/8KHvbtm3KOXnVX6NGlv9eqqioEPHVq1ctlpMfjd91110i9vRUf4TJK2rkVVCAutlhQx5mmjRpkohnzpwp4m7duinl8vPzRfznP/9ZxJZWKQHAzp07Rbx79+461ZPuLCAgQMSbN2+2eE7eeNG8Qa87rCTlExkiIiLSFjsyREREpC12ZIiIiEhbHoY8eGatoDRG7U6++OILEctzZCZMmKCUe/vtt+utTjIbm+eO6rP95HkV5vkR8vj666+/LuIRI0Yo5a5fvy5ieRfWpUuXKuU2bdpUp7o6mqu2X2xsrIizs7Mtljt+/Lhy/MEHH4h47NixIrY2b0JuP3OaBHlc/9SpUxavIf+fmj17tojlpd0A4OfnZ/EaaWlpIp46darFcjJ7tR/gvJ+h5p9lcnZ0SztXW2P+PuTPSJ4DZ+3/VX1y1XuwtuT5Z1u3bhXxAw88oJSTv++XX35ZxM8++6wDa2d/trQfn8gQERGRttiRISIiIm1xaIlDSy5h9OjRyrGlDc/MS3TlZdu///3vRewqm3y6avvJQwrp6enKuccff7zG1zN/nytXrhTxG2+8IWI53YE9JCYmKseZmZkWyxYXF4tYztxdXl5u8TXuMLQkZ+gFgObNmzvsveQl1+bMyuYl+vXFVe9Ba+TM2PKmjoD6e0r+/WW2bNkyEcvZe3XDoSUiIiJya+zIEBERkbY4tGRhaKlt27ZKOWsrKhxJx8ei9hASEiLiuLg4EcsbTQJqltczZ86IWN64EgAOHTpk7yraRIf2Cw8PV44PHz5c42uYv0/5Uba8MavZ559/LmJrG9bJw4vdu3cXsbyCCQD69u1r8Rr//ve/RTxy5EgRX7t2zeJrOLRUe+Ys6uPGjROxvAmvvNGkI+hwD5rJQ7/m4XR5413z0KpM3gT0hx9+sGPt6heHloiIiMitsSNDRERE2mJHhoiIiLTV4ObItGjRQjkuLCwUcVVVlYjbt2+vlLO2RNORdBzfdaTg4GDleO3atSKWd+E2724sz52w9xJga3RoPzlTKKDOL1qzZo1yLigoyK7vLd9/1uZKyO1u/j9gSWlpqXK8ZMkSEaemptp0Dc6RcQx5Z+0VK1Yo53Jzc0Vsj5+7OtyD1nTq1Ek5lueV+fr6itg81+s3v/mNiOXUA7rhHBkiIiJya+zIEBERkbY8nV2BmpKXpcmP1S5fvqyUMy9Zu8W8aZ38mHXGjBkidtZQEll38uRJ5Xjo0KEilocGW7durZSTj+tzaEkH5izIW7ZsEfGBAweUc3//+99FXJsMwGYRERF1voZMXmI9f/585dzBgwft+l66MG8SKH8uAQEBNl3jzTffFLH5HgwNDRXxtGnTRHz33XdbvN6AAQNEHB0drZzbvn27iOVl8sDtP+cbAvM9Iv/eky1fvlw51nk4qab4RIaIiIi0xY4MERERaUu7VUvy0ND06dNFbM5c+Oijj4pYzhgqz/gGgKZNm4pYzh7qKsMPOs64X7dunYjNKyZWrVol4uPHj9f42n5+fsrxrFmzRLx06VIRX7hwQSnXq1cvERcVFdX4fWtLx/azxtvbW8RyVt1BgwYp5T799FMRBwYGinjSpEkWy8l++9vfKsfykJacMXb16tVKOXnlk7WMvbZyh1VL9alVq1YilocoAXWjTk9P22Y1yKubAMvDg3/4wx8sXkPHe7Br164i/s9//qOck7OZy+TPF3BeNnN746olIiIicmvsyBAREZG22JEhIiIibbn8HBkvLy/luKCgQMTyOKLZ5s2bRdy5c2cRm7Mkyjv9RkVFidi8c6uz6Di+K7eRedy2srJSxPIurhcvXlTKyXMn5KWa8s7HgNpm8ryY4cOHK+U+/vhjm+pubzq2H/0f58jYj3xP9u7dW8Rz5sxRyjVqZPnv6xs3bohYXqYdHx9v8TU63oNxcXEiNs81kslpQsxZt81pFXTFOTJERETk1tiRISIiIm25/NCSmfwYcu7cuSL29/e3+Bq57ubMkCNGjBBxTk6OHWpoXzo+FpWXsT/11FPKOTl7sqVlhNZcv35dOT516pSIx44dK2JnDSWZ6dh+9H8cWnI887DQPffcY7HsRx99JOJ9+/bZdH0d78EePXqIOCMjQzkXHh4u4oSEBBHLQ23uhENLRERE5NbYkSEiIiJtsSNDRERE2tJujowsOTlZxPJyNQDo37+/iD/77DMRy2nsAetL21yBjuO71shL5s3LBW1h3tV87969da6TI7lb+zU0nCOjP96DeuMcGSIiInJr7MgQERGRtrQeWmoI+FhUb2w/vXFoSX+8B/XGoSUiIiJya+zIEBERkbbYkSEiIiJtsSNDRERE2mJHhoiIiLTFjgwRERFpix0ZIiIi0hY7MkRERKQtdmSIiIhIWzZn9iUiIiJyNXwiQ0RERNpiR4aIiIi0xY4MERERaYsdGSIiItIWOzJERESkLXZkiIiISFvsyBAREZG22JEhIiIibbEjQ0RERNr6H0knrLi2CapSAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 700x250 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "images, labels = next(iter(train_loader))\n",
    "idx = 0  # Change the index to see different images\n",
    "\n",
    "# Show some images\n",
    "fig, axes = plt.subplots(1, 5, figsize=(7, 2.5))\n",
    "for i in range(5):\n",
    "    axes[i].imshow(images[i+idx][0], cmap='gray')\n",
    "    axes[i].set_title(f\"label: {labels[i+idx]}\")\n",
    "    axes[i].axis('off')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Define the neural network structure\n",
    "\n",
    "For the MNIST task, it is sufficient to use a small, fully-connected neural network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleNN(nn.Module):\n",
    "    def __init__(self, input_size=28*28, num_classes=10, hidden_size=128):\n",
    "        super(SimpleNN, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.fc1 = nn.Linear(input_size, hidden_size)\n",
    "        self.activation = nn.ReLU()\n",
    "        self.fc2 = nn.Linear(hidden_size, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        x = self.fc1(x)\n",
    "        x = self.activation(x)\n",
    "        x = self.fc2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Create the NN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SimpleNN(\n",
      "  (flatten): Flatten(start_dim=1, end_dim=-1)\n",
      "  (fc1): Linear(in_features=784, out_features=128, bias=True)\n",
      "  (activation): ReLU()\n",
      "  (fc2): Linear(in_features=128, out_features=10, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = SimpleNN().to(device)\n",
    "print(model) # Look at the model architecture"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Define the loss function\n",
    "\n",
    "Here we are dealing with a classification problem with 10 classes (digits from 0 to 9), which loss function should we use?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "**Cross-entropy Loss**\n",
    "\n",
    "For one data pair $(x,y)$, the cross-entropy loss is calculated\n",
    "\n",
    "$$\n",
    "l_\\text{Cross-Entropy}(y,y') = - \\sum_{i=1}^{C} \\log \\frac{\\exp{y_i}}{\\sum_{c=1}^{C} \\exp{(y_{c})}} y'_{i},\n",
    "$$\n",
    "\n",
    "Here \n",
    "- $\\{1, \\dots, C=10 \\}$ class indices. \n",
    "- $y'$ is a one-hot vector, i.e. $y'_{i}=1$ if the ground truth label is class $i$. \n",
    "- $y$ is the output predicted by the neural network."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "$$\n",
    "l_\\text{Cross-Entropy}(y,y') = - \\sum_{i=1}^{C} \\log \\frac{\\exp{y_i}}{\\sum_{c=1}^{C} \\exp{(y_{c})}} y'_{i},\n",
    "$$\n",
    "\n",
    "- Note 1: The first part $\\exp{y_i}/\\sum_c\\exp{y_c}$ is a `Softmax` activation, mapping the unbounded outputs to probabilities between $[0,1]$.\n",
    "- Note 2: For the batched input, the loss is commonly averaged over the batches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the loss function\n",
    "criterion = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Define an optimizer\n",
    "\n",
    "We can choose from the common optimizers, see [`torch.optim` documentation](https://pytorch.org/docs/stable/optim.html):\n",
    "- Stochastic gradient descent ([SGD](https://pytorch.org/docs/stable/generated/torch.optim.SGD.html#torch.optim.SGD)): `optim.SGD`\n",
    "- [Adam](https://pytorch.org/docs/stable/generated/torch.optim.Adam.html#torch.optim.Adam): `optim.Adam`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Feel free to change the optimizer and its hyperparameters\n",
    "learning_rate = 1e-3\n",
    "optimizer = optim.Adam(model.parameters(), lr=learning_rate)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Define the training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, train_loader, criterion, optimizer, num_epochs=10):\n",
    "    model.train()  # Set the model to training mode\n",
    "    for epoch in range(num_epochs): # Loop through the dataset multiple epochs\n",
    "        for i, (images, labels) in enumerate(train_loader): # Loop through batches in the train loader\n",
    "            images, labels = images.to(device), labels.to(device) # Move data to device (GPU if available)\n",
    "            outputs = model(images) # Forward pass\n",
    "            loss = criterion(outputs, labels) # Compute loss\n",
    "\n",
    "            optimizer.zero_grad() # Prepare for backward pass\n",
    "            loss.backward() # Backward pass\n",
    "            optimizer.step() # Update model parameters\n",
    "\n",
    "            if (i+1) % 400 == 0:\n",
    "                print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item():.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the cell below to start training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/10], Step [400/938], Loss: 0.2974\n",
      "Epoch [1/10], Step [800/938], Loss: 0.4021\n",
      "Epoch [2/10], Step [400/938], Loss: 0.1189\n",
      "Epoch [2/10], Step [800/938], Loss: 0.1314\n",
      "Epoch [3/10], Step [400/938], Loss: 0.2653\n",
      "Epoch [3/10], Step [800/938], Loss: 0.1246\n",
      "Epoch [4/10], Step [400/938], Loss: 0.2013\n",
      "Epoch [4/10], Step [800/938], Loss: 0.0542\n",
      "Epoch [5/10], Step [400/938], Loss: 0.1582\n",
      "Epoch [5/10], Step [800/938], Loss: 0.1365\n",
      "Epoch [6/10], Step [400/938], Loss: 0.0242\n",
      "Epoch [6/10], Step [800/938], Loss: 0.0597\n",
      "Epoch [7/10], Step [400/938], Loss: 0.0699\n",
      "Epoch [7/10], Step [800/938], Loss: 0.0291\n",
      "Epoch [8/10], Step [400/938], Loss: 0.0966\n",
      "Epoch [8/10], Step [800/938], Loss: 0.0313\n",
      "Epoch [9/10], Step [400/938], Loss: 0.1028\n",
      "Epoch [9/10], Step [800/938], Loss: 0.0948\n",
      "Epoch [10/10], Step [400/938], Loss: 0.0491\n",
      "Epoch [10/10], Step [800/938], Loss: 0.0685\n"
     ]
    }
   ],
   "source": [
    "train(model, train_loader, criterion, optimizer, num_epochs=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Evaluating the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Accuracy: 96.73%\n"
     ]
    }
   ],
   "source": [
    "# Evaluation on the test set\n",
    "model.eval()\n",
    "correct = 0\n",
    "total = 0\n",
    "\n",
    "with torch.no_grad():\n",
    "    for images, labels in test_loader:\n",
    "        images, labels = images.to(device), labels.to(device)\n",
    "        outputs = model(images)\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "        total += labels.size(0)\n",
    "        correct += (predicted == labels).sum().item()\n",
    "\n",
    "accuracy = correct / total\n",
    "print(f\"Test Accuracy: {accuracy * 100:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Visualize the model predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAjIAAACcCAYAAABlRUHuAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/H5lhTAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAgQElEQVR4nO3de1hU1foH8O+ggoCXgFBEEH0gRFHTNI+iiMcLaSZqeUw7WqGlZqTmSTStUJNMLcMyb2VeC8W0NJ/Kk4nlJStBvGSEmpCaeb8EEiis3x/9XGftkcuADDNr+H6ex+d5995r9rzO24yrvfZa2ySEECAiIiLSkJOtEyAiIiIqL3ZkiIiISFvsyBAREZG22JEhIiIibbEjQ0RERNpiR4aIiIi0xY4MERERaYsdGSIiItIWOzJERESkLat3ZFasWAGTyYTMzMwyv7Zr165o0aJFhebTuHFjPPnkkxV6TkfHGuqN9dMb66c31s/6eEWmjM6cOYORI0eiSZMmcHV1RWBgICZMmICLFy/aOjWyUHx8PKKiolC/fn2YTCZMmzbN1imRhTIzM2EymYr8s3btWlunRxY4duwYBg4cCA8PD7i5uaFz585ITk62dVpkIXv8/axu6wR0kp2djY4dOyInJwdjxoyBv78/Dhw4gAULFiA5ORkpKSlwcmLf0N699NJL8PHxQZs2bbB161Zbp0PlMGTIEDz44IOGfR07drRRNmSpkydPomPHjqhWrRomTpwId3d3LF++HJGRkfj666/RpUsXW6dIpbDH3092ZMpg8+bNyMrKwpYtW9CnTx+539PTEzNmzMCBAwfQpk0bG2ZIljhx4gQaN26MCxcuwNvb29bpUDncd999GDp0qK3ToDJ6/fXXceXKFRw+fBhNmzYFADz99NMICQnB888/j5SUFBtnSKWxx99Pm1w+2LRpE/r06QNfX1+4uLggMDAQr776KgoKCopsn5KSgrCwMLi6uqJJkyZYvHjxbW3y8vIQFxeHoKAguLi4wN/fH7GxscjLyys1n+PHj+P48eOltrt27RoAoH79+ob9DRo0AAC4urqWeg5HoWsNgb/HiKs6net3S05ODvLz88v0Gkeha/127tyJNm3ayE4MALi5uSEqKgqpqak4evRoqedwBLrWD7DP30+bXJFZsWIFatWqhQkTJqBWrVrYvn07XnnlFVy7dg1z5841tL18+TIefPBBDBo0CEOGDEFSUhKeeeYZODs7Y/jw4QCAwsJCREVFYdeuXRg5ciSaNWuGQ4cO4a233kJGRgY+/fTTEvPp3r07AJR6M1aXLl3g5OSEcePG4c0334Sfnx8OHjyI+Ph49O/fHyEhIeX+THSjaw3pb7rXb/r06Zg4cSJMJhPatm2L+Ph4REZGlvlz0JWu9cvLy4OHh8dt+93c3AD8/Q/2PffcY+GnoC9d62e3hJUtX75cABAnTpyQ+65fv35bu1GjRgk3Nzfx119/yX0RERECgHjzzTflvry8PNG6dWtRr149kZ+fL4QQYvXq1cLJyUns3LnTcM7FixcLAGL37t1yX0BAgHjiiScM7QICAkRAQIBFf5/3339f3HXXXQKA/PPEE0+IGzduWPR6HTlaDW85f/68ACDi4uLK9DrdOFL9srKyRGRkpFi0aJHYvHmzSEhIEI0aNRJOTk5iy5Ytpb5eR45Uv759+4q77rpLXLt2zbC/Y8eOAoB44403Sj2Hbhypfip7+v20ydCSOgTz559/4sKFCwgPD8f169eRnp5uaFu9enWMGjVKbjs7O2PUqFE4d+6cHE9dv349mjVrhpCQEFy4cEH+6datGwCUekd8ZmamxT3Rhg0bon379khISMAnn3yCCRMm4MMPP8TkyZMter2j0LmGpG/9GjVqhK1bt2L06NHo27cvxo0bh/3798Pb2xv/+c9/LP3ra0/X+j3zzDO4cuUKHn30Uezfvx8ZGRkYP3489u3bBwDIzc216O+vO13rZ69sMrT0008/4aWXXsL27dvlfSe3XL161bDt6+sLd3d3w77g4GAAf3/4HTp0wNGjR/Hzzz8Xe+PRuXPnKiTv3bt346GHHsLevXvRrl07AED//v1Rp04dTJ8+HcOHD0fz5s0r5L3sna41pL85Uv08PT0RHR2N119/HadOnYKfn5/V3ste6Fq/3r1745133sHkyZNx3333AQCCgoIQHx+P2NhY1KpVq0Lex97pWj97VekdmStXriAiIgJ16tTBjBkzEBgYiJo1ayI1NRWTJk1CYWFhmc9ZWFiIli1bYt68eUUe9/f3v9O0AQBLlixB/fr1ZSfmlqioKEybNg179uypEh0ZnWtIjlm/W+e/dOmSw3dkdK9fTEwMoqOjcfDgQTg7O6N169ZYtmwZgP/9A+3IdK+fPar0jsyOHTtw8eJFbNy40bBmwIkTJ4ps//vvvyMnJ8fQI83IyADwv7unAwMDceDAAXTv3h0mk8lquZ89e7bIu8pv3LgBALh586bV3tue6FxDcsz6/frrrwBgN9NBrckR6ufu7m5Y92fbtm1wdXVFp06drP7etuYI9bM3lX6PTLVq1QAAQgi5Lz8/HwsXLiyy/c2bN7FkyRJD2yVLlsDb2xtt27YFAAwaNAinT5/Ge++9d9vrc3NzkZOTU2JOlk49Cw4OxtmzZ7Fjxw7D/sTERACoMmvI6FxD0rt+58+fv23f6dOn8cEHH6BVq1ZyKQRHpnP9irJnzx5s3LgRI0aMQN26dct1Dp04Wv3sQaVfkQkLC4OHhweeeOIJjB07FiaTCatXrzYUVeXr64vZs2cjMzMTwcHBWLduHdLS0rB06VLUqFEDADBs2DAkJSVh9OjRSE5ORqdOnVBQUID09HQkJSVh69attw0HqSydehYTE4Ply5ejb9++eO655xAQEIBvvvkGiYmJ6NmzJ/7xj3+U70PRjM41BIDVq1cjKysL169fBwB8++23mDlzpswjICCgLB+HdnSuX2xsLI4fP47u3bvD19cXmZmZWLJkCXJycjB//vzyfSCa0bl+WVlZGDRoEKKiouDj44OffvoJixcvRqtWrfDaa6+V7wPRjM71A+z099Pa06KKmnq2e/du0aFDB+Hq6ip8fX1FbGys2Lp1qwAgkpOTZbuIiAgRGhoq9u3bJzp27Chq1qwpAgICxIIFC257n/z8fDF79mwRGhoqXFxchIeHh2jbtq2YPn26uHr1qmx3p1PP0tPTxcCBA4W/v7+oUaOGCAgIEC+88ILIyckpy8eiFUer4a0pjUX9UXN3FI5Uv48++kh06dJFeHt7i+rVq4u7775bDBgwQKSkpJT1Y9GGI9Xv0qVLol+/fsLHx0c4OzuLJk2aiEmTJt02HduROFL9buVkb7+fJiGK6QYSERER2Tk+4ZCIiIi0xY4MERERaYsdGSIiItIWOzJERESkLXZkiIiISFvsyBAREZG22JEBYDKZMG3aNFunQeXE+umN9dMb66c3R6gfOzIV4OjRoxg8eDD8/Pzg5uaGkJAQzJgxQ658SPYrOzsbcXFx6NWrFzw9PWEymbBixQpbp0UW2rFjB0wmU5F/9u7da+v0yAIpKSno1asX6tSpg9q1ayMyMhJpaWm2TossYC+/n5X+iAJHc/LkSbRv3x5169ZFTEwMPD098d133yEuLg4pKSnYtGmTrVOkEly4cAEzZsxAo0aNcO+99972HC3Sw9ixY3H//fcb9gUFBdkoG7JUamoqOnfuDH9/f8TFxaGwsBALFy5EREQEfvjhBzRt2tTWKVIJ7OX3U+uOzPXr1+Hm5mbTHFavXo0rV65g165dCA0NBQCMHDkShYWFWLVqFS5fvgwPDw+b5miv7KF+DRo0wJkzZ+Dj44N9+/bd9o8hFc8e6ndLeHg4Bg4caOs0tGIP9Xv55Zfh6uqK7777Dl5eXgCAoUOHIjg4GFOmTMGGDRtsmp89s4f62cvvp02Hlm5dFl63bh2mTJkCHx8fuLu7IyoqCidPnjS07dq1K1q0aIGUlBR06dIFbm5umDJlCgAgLy8PcXFxCAoKgouLC/z9/REbG4u8vDzDOfLy8vD888/D29sbtWvXRlRUFE6dOlVkbunp6fjtt99K/Ttcu3YNAFC/fn3D/gYNGsDJyQnOzs4Wfx66cYT6ubi4wMfHp5yfgN4coX6qP//8Ezdv3izTa3TmCPXbuXMnevToITsxwN+/nREREdiyZQuys7PL+rFowxHqZy+/n3ZxRSY+Ph4mkwmTJk3CuXPnkJCQgB49eiAtLQ2urq6y3cWLF9G7d28MHjwYQ4cORf369VFYWIioqCjs2rULI0eORLNmzXDo0CG89dZbyMjIwKeffipf/9RTT2HNmjV47LHHEBYWhu3bt6NPnz5F5tSsWTNERESUeqmsa9eumD17NkaMGIHp06fDy8sLe/bswaJFizB27Fi4u7tXxEdk13SuHzlG/aKjo5GdnY1q1aohPDwcc+fOLfFpv45E5/rl5eUZcrzFzc0N+fn5OHz4MDp06FCuz0UXOtfPbtjkUZX/Lzk5WQAQDRs2NDz9NCkpSQAQ8+fPl/tuPXFz8eLFhnOsXr1aODk5iZ07dxr2L168WAAQu3fvFkIIkZaWJgCIMWPGGNo99thjAoCIi4sz7AcgIiIiLPp7vPrqq8LV1dXwFNCpU6da9FqdOUr9bvnxxx8FALF8+fIyvU5XjlC/3bt3i0ceeUQsW7ZMbNq0ScyaNUt4eXmJmjVritTUVEs+Bm05Qv1atmwpgoODxc2bN+W+vLw80ahRIwFAfPzxx6WeQ1eOUD+VLX8/7WLW0uOPP47atWvL7YEDB6JBgwb4/PPPDe1cXFwQHR1t2Ld+/Xo0a9YMISEhuHDhgvzTrVs3AEBycjIAyHONHTvW8Prx48cXmZMQwuLeaOPGjdGlSxcsXboUGzZswPDhw/Haa69hwYIFFr1ed7rXr6rTuX5hYWH4+OOPMXz4cERFRWHy5MnYu3cvTCYTXnzxxVJf7wh0rt+YMWOQkZGBESNG4MiRIzh8+DAef/xxnDlzBgCQm5tb6jl0p3P97IVdDC3dc889hm2TyYSgoCBkZmYa9jds2PC2e06OHj2Kn3/+Gd7e3kWe+9y5cwCArKwsODk5ITAw0HD8Tu+KX7t2LUaOHImMjAz4+fkBAB5++GEUFhZi0qRJGDJkiGH81xHpXD9yvPoFBQWhX79+2LhxIwoKClCtWrUKfw97onP9Ro8ejZMnT2Lu3LlYuXIlAKBdu3aIjY1FfHw8atWqdUfn14HO9bMXdtGRsVRRY6mFhYVo2bIl5s2bV+Rr/P39rZrTwoUL0aZNG9mJuSUqKgorVqzA/v370aNHD6vmoAt7rB9ZTqf6+fv7Iz8/Hzk5OahTp45NcrA39lq/+Ph4vPDCC/jpp59Qt25dtGzZUt7IGhwcbPX314W91s8e2EVH5ujRo4ZtIQSOHTuGVq1alfrawMBAHDhwAN27d4fJZCq2XUBAAAoLC3H8+HFDL/SXX34pf+IAzp49W+T06hs3bgBAlZhFoXP9yDHr9+uvv6JmzZpV4v/oHaF+Hh4e6Ny5s9zetm0b/Pz8EBISUiHnt2eOUD9bs4t7ZFatWoU///xTbn/88cc4c+YMevfuXeprBw0ahNOnT+O999677Vhubi5ycnIAQJ7r7bffNrRJSEgo8ryWTj8LDg7G/v37kZGRYdifmJgIJycni/5j1J3O9SO963f+/Pnb9h04cACbN29GZGQknJzs4ifOqnSuX1HWrVuHH3/8EePHj2f9SmGP9bMFu7gi4+npic6dOyM6Ohpnz55FQkICgoKC8PTTT5f62mHDhiEpKQmjR49GcnIyOnXqhIKCAqSnpyMpKQlbt25Fu3bt0Lp1awwZMgQLFy7E1atXERYWhq+//hrHjh0r8ryWTj+bOHEivvjiC4SHhyMmJgZeXl7YsmULvvjiCzz11FPw9fUtz0eiFZ3rBwALFizAlStX8PvvvwMAPvvsM7m+wnPPPYe6deta/mFoSOf6Pfroo3B1dUVYWBjq1auHI0eOYOnSpXBzc8Prr79eno9DOzrX79tvv8WMGTMQGRkJLy8v7N27F8uXL0evXr0wbty48nwc2tG5foCd/H5W+jwpxa3pZ4mJieLFF18U9erVE66urqJPnz4iKyvL0DYiIkKEhoYWeZ78/Hwxe/ZsERoaKlxcXISHh4do27atmD59urh69apsl5ubK8aOHSu8vLyEu7u76Nu3rzh58uQdTz/7/vvvRe/evYWPj4+oUaOGCA4OFvHx8eLGjRtl+jx04yj1CwgIMEydV/+cOHGiLB+JVhyhfvPnzxft27cXnp6eonr16qJBgwZi6NCh4ujRo2X+PHTjCPU7duyYiIyMFHfffbdwcXERISEhYtasWSIvL6/Mn4duHKF+QtjH76dddGTWr19vyzSonFg/vbF+emP99Mb6VRzHH4AkIiIih8WODBEREWmLHRkiIiLSlkkIIWydBBEREVF58IoMERERaYsdGSIiItKWxQvilbT8MVlPRY38sX62wfrprSJH3llD2+B3UG+W1I9XZIiIiEhb7MgQERGRttiRISIiIm2xI0NERETaYkeGiIiItMWODBEREWmLHRkiIiLSFjsyREREpC12ZIiIiEhb7MgQERGRttiRISIiIm2xI0NERETaYkeGiIiItMWODBEREWmLHRkiIiLSFjsyREREpK3qtk6gMgwYMEDGM2bMMBwLDQ2V8Q8//CDjqVOnGtp9/fXXVsqOrKFXr16G7cmTJ8t4woQJMk5NTa20nIgcnfpbq/62AsCOHTtkvGvXrspKicrg888/l/EDDzxgOHb48GEZd+3aVcaXL1+2el6l4RUZIiIi0hY7MkRERKQthx1aat26tYw3btwoY/XyGGAcgvjtt99k/PLLLxvatW3bVsZz5sypqDSpAj300EMyXrNmjeFYzZo1ZVy3bt1Ky6mqML8MfeDAARn/8ccfFp3DZDLJuKQatWnTRsbq9xIAOnXqJOOtW7cWe0zN99FHHzW0S05OtijfqqR+/foydnV1NRxbvHixjNUahoeHG9pNmzZNxurw7ttvv11RadIdEkIUGQPAkSNHZGwPw0kqXpEhIiIibbEjQ0RERNpy2KGluLg4GX/55Zcyfvjhhw3tcnNzi3z9448/btgeP368jGNjY2XMYSbbUmcjqXHt2rUN7dQhiLS0NKvnVRVERkbKWP2OAcCqVatkPHHiRBk3atTI0K5z584ybtKkiYyfe+65cuV08OBBGb/22muGY4mJiTJu3ry5jC9cuFCu93I0tWrVMmwvWLBAxl26dJGxh4eHod1dd90lY3U4Qh1mMnfPPfeUN02tOTs7FxlnZ2fbIh0AQEBAgIxbtGhhszzuBK/IEBERkbbYkSEiIiJtsSNDRERE2nKYe2RGjRpl2FbHvZ9++ukyn6+goMCwnZCQIOOvvvpKxj/++KOhHaduVi51pebCwkIZP/XUU4Z26nRgKr97771XxitXriy23bBhw2QcFRUl45KmVav3VJhP71y7dq2MO3bsKONPP/3U0E79nnp7exuOnT59WsbF3RtXlZnfl6QuZ6DeP2P+m6euBvvJJ5/IuF+/foZ26v0zap2qEh8fHxkHBwfLeNu2bbZIBwDQrFkzGTds2LDYdps2baqMdMqFV2SIiIhIW+zIEBERkba0HlpSp0j379/fcMz8suadUoea1EurXBXU+tTpnrNmzTIcq1GjhozV1XyXL19u/cSqgMaNGxu21Sm56mqvJSlpOEl9eOCJEydkrE6VBoDt27fLWF2l+dq1a8We++rVqxblV9WEhITIOD09XcbHjx83tDt27JiMFy1aJOMNGzYY2hU3dVg9N/1NXT1ejSuT+fDRsmXLimy3b98+w7Y6hGhveEWGiIiItMWODBEREWmLHRkiIiLSlnb3yLi5ucl43LhxMn722WcN7fLz8yslH0vvE6CyUZc9T0pKknG3bt0M7dR7ksaOHWv1vKqaJ5980rCtPkFadePGDcP2li1bZLxjxw4Zr1+/3tBOfTK2+dN2i1NZ321Hof5mAsDdd98t46VLl8p48ODBhnbqIx7U+2LM74lxcXEp8n3z8vLKnixZnflyJMX9G3b+/HnDdkn3o9kar8gQERGRttiRISIiIm1pPbSkPkH15MmTlfa+6pNg1ScuU8Xx8/OTsflwkkq95G2+GizduaysrGKPTZ06Vcbm0+LJtmJiYmT8yiuvGI6pQ0vqaq3qfqD4IbygoCDDtjpcWK9ePRkPHTrU0I5LU9iO+pT5MWPGFNtOXRF/7ty5Vs2pIvGKDBEREWmLHRkiIiLSlnZDSyr1QWZqXBHM78RXH1p36tQpGX/zzTcV+r70t549exa5f926dYbtd999tzLSqbLMVx+9efOmjJs3b17Z6ZCFevToIWNPT0/DMXVofM+ePTJWH7pqLjAwUMbqgyEBIDQ0tMjXTJkyxbDNoSXb+de//iVj8/8eVM8//7yMv/32W6vmVJF4RYaIiIi0xY4MERERaYsdGSIiItKWSVi4nKbJZLJ2LmX20Ucfydj8acdfffVVmc9Xu3ZtGR86dMhwTH3yZ0nT1yqapaudlsYe66dydXU1bGdkZMjY19dXxupYPQBkZmZaNa875Wj1S0tLk3GrVq1kHB0dbWi3cuXKykrJqiqqfkDl1lB9Qrj5d6u4ZQrq1Klj2FZX+lXvnWjatGmx73vp0iUZR0ZGGo6lpqaWkLH1ONp30FJdu3aV8ebNm2VsvtKzuoLvP//5Txnby9PLLakfr8gQERGRttiRISIiIm1pPf1avczdunVrwzFLh5YeeughGc+ZM0fG5qtS7tq1q+wJksVGjBhh2FaHk9544w0Znz59utJyotu9//77Mp43b56MP/jgA0O7ixcvylh9gCRZhzqMAADTpk2Tsfr9AYwP//Px8ZHxDz/8YGin1lRdXbukoSV1CMNWQ0n0N3VFZ/PhJNWECRNkbC/DSWXFKzJERESkLXZkiIiISFtaDy2pd8ibz5pISEiQsbqS4cyZMw3t1EurDz74oIztfTaMI1BXY+7fv7/h2B9//CHjBQsWyPjGjRtWz4uKp9ZiwIABMlZnOwDGy9o7d+6U8dWrV62YXdWlrtYLAB06dJDxpEmTDMeWLVsm423btsk4NzfX0E5dxbkk6orA6nAjVS5nZ2fDtvqdVGt05MgRQ7vExETrJlYJeEWGiIiItMWODBEREWmLHRkiIiLSltYr+/r5+cnYfNxvyZIlMlbvfXnnnXcM7d577z0ZFxQUVHSKd8yRV6WMjY2V8axZswzHDh8+LON777230nKqaI5cvwYNGsh47969hmP+/v4ybt++vYz37dtn/cQqkK4r+6qrLv/111+GY+rUeDUuibpK7/Tp0w3H7r//fhmHhITI+NixY5Yla2WO/B10d3eXsfkyB+p9U+q/bcOGDTO0W7dunZWyqxhc2ZeIiIgcGjsyREREpC2tp1+rU3S//PJLw7F///vfMm7btq2Mz5w5Y/3EyCLq9F1z6tAS2Sf1u9SvXz/DMXWFV3X4NiIiwtBOXWWWKs7Bgwfv+ByPPPKIjIsbPgKMU3uzs7Pv+H3Jci+88IKMw8PDi22nPlTZ3oeSyoNXZIiIiEhb7MgQERGRtrQeWoqJiZFxnz59DMfy8/NlrN7ZTfbJfMVe9QGeVLnq1q0rY0tX4lUf4AoYZ6GpDy0MCwsztDMfEibb6dGjh2F7xYoVMi7pNzQ+Pl7G6nA/WYeXl5eMn3322WLbZWVlyXju3LlWzcnWeEWGiIiItMWODBEREWmLHRkiIiLSlnYr+3bu3FnGa9askbG6ei8A9O7dW8bqCpObNm2yYnYVz5FXpfzuu+9kHBoaajhWp06dyk7HKnSs35AhQ2T8+eefG45Zes+Mm5ubjA8dOiRj8xW4+/btW54UK42uK/uWh6+vr2FbnUJ/3333yfjXX381tOvatauMT506ZZ3k7oCO38GSfPHFFzLu2bNnse2aNm0q4+PHj1s1J2viyr5ERETk0NiRISIiIm1pN/1avVy2Z88eGZtfsq5du7aMFy5cKOOdO3ca2l26dKmiUyTSWlBQkIzVqdiA5UNL169fl7G6ei+XQrBf6kN4AaB169ZFtlOHNgD7HE5yJOrDWQGgRYsWRbYzHz7SeTiprHhFhoiIiLTFjgwRERFpS7uhpV69esm4pEtn33//vYxv3rwp4wceeMDQLjExsQKzo7JQV/NVZ7kAxod+fvjhh5WWU1VkPhtDnaHy7rvvluucHTp0kHGTJk1kfOXKlXKdj6xPfbguADg5/e//c9W6lfe/CSofdaYucPtQ0y2vvvpqZaRjl3hFhoiIiLTFjgwRERFpix0ZIiIi0pZ298isX79exuHh4Ra95syZMzLmPTL2Izo6WsapqamGY+ZTQcl61HtiAKBfv34yNn9qrrrkgapdu3aG7VWrVslYXaX5v//9b7nzpIqnroA+ZcqUYtutXbtWxunp6VbNiSynfs8++ugjG2ZiW7wiQ0RERNpiR4aIiIi0pd3Q0pYtW2Q8fvx4GT/77LOGdmlpaTJu3ry5jJctW2a13Khs1Onz5lOsp06dKmP1oZ8bNmywfmJVzNmzZw3b2dnZMp43b57h2OjRo2Xcpk0bGb/11luGdupw0i+//CLjRYsW3VmyVKEGDx4s44YNGxqOZWRkyHjOnDmVlhNZTl1pu7Cw0IaZ2BavyBAREZG22JEhIiIibbEjQ0RERNrS7h4Z9am66jTR2NhYQ7uhQ4fKeOnSpTKeP3++FbOj8vrss88M26NGjZLxihUrZJyfn1/i66jszJ9erH7eMTExhmPm0+RvKSgoMGyr96LNnDlTxnxEgT5Wrlwp48zMTNslUsX9/PPPhu3Lly/bKBP7xSsyREREpC12ZIiIiEhbJiGEsKih2RNyqXJYWJ5SsX62oWP9XFxcZDxhwgTDMXX11x07dsh44sSJhnaOsvprRdUP4HfQVnT8DtL/WFI/XpEhIiIibbEjQ0RERNri0JKd42VRvbF+euPQkv74HdQbh5aIiIjIobEjQ0RERNpiR4aIiIi0xY4MERERaYsdGSIiItIWOzJERESkLYunXxMRERHZG16RISIiIm2xI0NERETaYkeGiIiItMWODBEREWmLHRkiIiLSFjsyREREpC12ZIiIiEhb7MgQERGRttiRISIiIm39H+CWAVPFFG3cAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 700x250 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Redefine a loader to allow for visualization of random test images\n",
    "random_test_loader = DataLoader(dataset=test_dataset, batch_size=5, shuffle=True)\n",
    "\n",
    "images, labels = next(iter(random_test_loader))\n",
    "images, labels = images.to(device), labels.to(device)\n",
    "outputs = model(images)\n",
    "_, predictions = torch.max(outputs, 1)\n",
    "\n",
    "# Show some images\n",
    "fig, axes = plt.subplots(1, 5, figsize=(7, 2.5))\n",
    "for i in range(5):\n",
    "    axes[i].imshow(images[i][0], cmap='gray')\n",
    "    axes[i].set_title(f\"label: {labels[i]} \\n pred: {predictions[i]}\")\n",
    "    axes[i].axis('off')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### What's next\n",
    "\n",
    "Change the hyperparameters of the network, training process, and observe what will happen"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bonus\n",
    "\n",
    "Try a different network structure, which one would you use?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "new_rl",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}