example_network.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Task\n",
    "Our task is simple, recognize handwritten digits. We will use MNIST dataset for this tutorial.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import necessary library\n",
    "In this tutorial, we are going to use pytorch, the cutting-edge deep learning framework to complete our task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from tqdm.notebook import tqdm\n",
    "from torchsummary import summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Create dataloader, in PyTorch, we feed the trainer data with use of dataloader\n",
    "## We create dataloader with dataset from torchvision, \n",
    "## and we dont have to download it seperately, all automatically done\n",
    "\n",
    "# Define batch size, batch size is how much data you feed for training in one iteration\n",
    "batch_size_train = 64 # We use a small batch size here for training\n",
    "batch_size_test = 1024 #\n",
    "\n",
    "# define how image transformed\n",
    "image_transform = torchvision.transforms.Compose([\n",
    "                               torchvision.transforms.ToTensor(),\n",
    "                               torchvision.transforms.Normalize(\n",
    "                                 (0.1307,), (0.3081,))\n",
    "                             ])\n",
    "#image datasets\n",
    "train_dataset = torchvision.datasets.MNIST('dataset/', \n",
    "                                           train=True, \n",
    "                                           download=True,\n",
    "                                           transform=image_transform)\n",
    "test_dataset = torchvision.datasets.MNIST('dataset/', \n",
    "                                          train=False, \n",
    "                                          download=True,\n",
    "                                          transform=image_transform)\n",
    "#data loaders\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset,\n",
    "                                           batch_size=batch_size_train, \n",
    "                                           shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(test_dataset,\n",
    "                                          batch_size=batch_size_test, \n",
    "                                          shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Label: tensor(7)\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAaEAAAGdCAYAAAC7EMwUAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAYuklEQVR4nO3df2jU9x3H8df564zuchBscpeZhiDKSiNC1anBanQzM2wym5VZu43YP8S20VXSInMyDPvDFDelf2R1TIrTVVf3h7WCUk3RJBbnloql4oqkGGc2E4KhvYupXqZ+9kfwtms09Xve5Z3LPR9w4N19P97bb7/N02/u8tXnnHMCAMDAGOsBAADZiwgBAMwQIQCAGSIEADBDhAAAZogQAMAMEQIAmCFCAAAz46wH+Kq7d+/q2rVrCgQC8vl81uMAADxyzqm3t1eFhYUaM2boc50RF6Fr166pqKjIegwAwCPq6OjQ1KlTh9xmxH07LhAIWI8AAEiBh/l6nrYIvfnmmyopKdHEiRM1e/ZsnT59+qHW8S04ABgdHubreVoidPDgQW3cuFFbtmzR+fPn9fTTT6uyslJXr15Nx8sBADKULx1X0Z43b56eeuop7dq1K/7YE088oZUrV6q+vn7ItdFoVMFgMNUjAQCGWSQSUW5u7pDbpPxMqL+/X+fOnVNFRUXC4xUVFTpz5syg7WOxmKLRaMINAJAdUh6h69ev686dOyooKEh4vKCgQF1dXYO2r6+vVzAYjN/4ZBwAZI+0fTDhq29IOefu+ybV5s2bFYlE4reOjo50jQQAGGFS/nNCU6ZM0dixYwed9XR3dw86O5Ikv98vv9+f6jEAABkg5WdCEyZM0OzZs9XY2JjweGNjo8rKylL9cgCADJaWKybU1tbqZz/7mebMmaMFCxboD3/4g65evaoXX3wxHS8HAMhQaYnQqlWr1NPTo1//+tfq7OxUaWmpjh07puLi4nS8HAAgQ6Xl54QeBT8nBACjg8nPCQEA8LCIEADADBECAJghQgAAM0QIAGCGCAEAzBAhAIAZIgQAMEOEAABmiBAAwAwRAgCYIUIAADNECABghggBAMwQIQCAGSIEADBDhAAAZogQAMAMEQIAmCFCAAAzRAgAYIYIAQDMECEAgBkiBAAwQ4QAAGaIEADADBECAJghQgAAM0QIAGCGCAEAzBAhAIAZIgQAMEOEAABmiBAAwAwRAgCYIUIAADNECABghggBAMwQIQCAGSIEADBDhAAAZogQAMAMEQIAmCFCAAAzRAgAYIYIAQDMECEAgBkiBAAwQ4QAAGaIEADADBECAJghQgAAM0QIAGCGCAEAzBAhAICZlEeorq5OPp8v4RYKhVL9MgCAUWBcOn7TJ598Uh988EH8/tixY9PxMgCADJeWCI0bN46zHwDA10rLe0JtbW0qLCxUSUmJnnvuOV2+fPmB28ZiMUWj0YQbACA7pDxC8+bN0759+3T8+HHt3r1bXV1dKisrU09Pz323r6+vVzAYjN+KiopSPRIAYITyOedcOl+gr69P06ZN06ZNm1RbWzvo+VgsplgsFr8fjUYJEQCMApFIRLm5uUNuk5b3hP7f5MmTNXPmTLW1td33eb/fL7/fn+4xAAAjUNp/TigWi+nTTz9VOBxO90sBADJMyiP02muvqbm5We3t7frb3/6mZ599VtFoVNXV1al+KQBAhkv5t+P+9a9/afXq1bp+/boee+wxzZ8/X2fPnlVxcXGqXwoAkOHS/sEEr6LRqILBoPUYAIBH9DAfTODacQAAM0QIAGCGCAEAzBAhAIAZIgQAMEOEAABmiBAAwAwRAgCYIUIAADNECABghggBAMwQIQCAGSIEADBDhAAAZogQAMAMEQIAmCFCAAAzRAgAYIYIAQDMECEAgBkiBAAwQ4QAAGaIEADADBECAJghQgAAM0QIAGBmnPUAlvr7+5Na19LS4nnNiRMnknqtkezSpUue17z33ntpmGSwn/zkJ0mtO336tOc1V69eTeq1AHAmBAAwRIQAAGaIEADADBECAJghQgAAM0QIAGCGCAEAzBAhAIAZIgQAMEOEAABmiBAAwAwRAgCY8TnnnPUQ/y8ajSoYDA7LayX7Rx9hu8zM7du3Pa+5detWGiYZbNKkSUmti8VintfcuXMnqdcayXp7ez2v+c1vfuN5TTL7e9euXZ7XwEYkElFubu6Q23AmBAAwQ4QAAGaIEADADBECAJghQgAAM0QIAGCGCAEAzBAhAIAZIgQAMEOEAABmiBAAwAwRAgCYGWc9ADLXuHHeD59vfOMbaZgkdXJycqxHGBGS+e+0c+dOz2uSuRjw2LFjPa9paGjwvAbDgzMhAIAZIgQAMOM5Qi0tLVqxYoUKCwvl8/l0+PDhhOedc6qrq1NhYaFycnJUXl6uixcvpmpeAMAo4jlCfX19mjVr1gO/x7p9+3bt3LlTDQ0Nam1tVSgU0rJly5L6R7IAAKOb53eWKysrVVlZed/nnHN64403tGXLFlVVVUmS9u7dq4KCAh04cEDr1q17tGkBAKNKSt8Tam9vV1dXlyoqKuKP+f1+LV68WGfOnLnvmlgspmg0mnADAGSHlEaoq6tLklRQUJDweEFBQfy5r6qvr1cwGIzfioqKUjkSAGAES8un43w+X8J959ygx+7ZvHmzIpFI/NbR0ZGOkQAAI1BKf1g1FApJGjgjCofD8ce7u7sHnR3d4/f75ff7UzkGACBDpPRMqKSkRKFQSI2NjfHH+vv71dzcrLKyslS+FABgFPB8JnTjxg199tln8fvt7e36+OOPlZeXp8cff1wbN27Utm3bNH36dE2fPl3btm3TpEmT9Pzzz6d0cABA5vMcoY8++khLliyJ36+trZUkVVdX649//KM2bdqkmzdv6uWXX9bnn3+uefPm6cSJEwoEAqmbGgAwKvhcMlcQTKNoNKpgMDgsr7V69eqk1v3pT3/yvGbMGO/f+XzllVc8r2lvb/e8ZuXKlZ7XjFYPeu9yKN///vfTMAke5ODBg57XJPv/Oh5NJBJRbm7ukNtw7TgAgBkiBAAwQ4QAAGaIEADADBECAJghQgAAM0QIAGCGCAEAzBAhAIAZIgQAMEOEAABmiBAAwAwRAgCYyeqraCfrpz/9qec13/ve9zyveemllzyvuXHjhuc1+J+JEyd6XnPvXxQeiZ599tmk1q1bt87zmmnTpiX1Wl4l8yUrmSvSS1JDQ0NS6zCAq2gDAEY0IgQAMEOEAABmiBAAwAwRAgCYIUIAADNECABghggBAMwQIQCAGSIEADBDhAAAZogQAMDMOOsBMtHbb789LGsw/G7duuV5zZUrV1I/SIr89re/TWrdE0884XnNcF3A1OfzeV4z0i+KnM04EwIAmCFCAAAzRAgAYIYIAQDMECEAgBkiBAAwQ4QAAGaIEADADBECAJghQgAAM0QIAGCGCAEAzHABUwCDfPDBB57XvPDCC2mYZLD//Oc/ntfs3r07DZMgFTgTAgCYIUIAADNECABghggBAMwQIQCAGSIEADBDhAAAZogQAMAMEQIAmCFCAAAzRAgAYIYIAQDMcAFTAIN897vftR7hgZxzntd0d3enYRKkAmdCAAAzRAgAYMZzhFpaWrRixQoVFhbK5/Pp8OHDCc+vWbNGPp8v4TZ//vxUzQsAGEU8R6ivr0+zZs1SQ0PDA7dZvny5Ojs747djx4490pAAgNHJ8wcTKisrVVlZOeQ2fr9foVAo6aEAANkhLe8JNTU1KT8/XzNmzNDatWuH/GRKLBZTNBpNuAEAskPKI1RZWan9+/fr5MmT2rFjh1pbW7V06VLFYrH7bl9fX69gMBi/FRUVpXokAMAIlfKfE1q1alX816WlpZozZ46Ki4t19OhRVVVVDdp+8+bNqq2tjd+PRqOECACyRNp/WDUcDqu4uFhtbW33fd7v98vv96d7DADACJT2nxPq6elRR0eHwuFwul8KAJBhPJ8J3bhxQ5999ln8fnt7uz7++GPl5eUpLy9PdXV1+tGPfqRwOKwrV67ol7/8paZMmaJnnnkmpYMDADKf5wh99NFHWrJkSfz+vfdzqqurtWvXLl24cEH79u3TF198oXA4rCVLlujgwYMKBAKpmxoAMCp4jlB5efmQFxA8fvz4Iw0EAMgeXDsOAGCGCAEAzBAhAIAZIgQAMEOEAABmiBAAwAwRAgCYIUIAADNECABghggBAMwQIQCAGSIEADBDhAAAZogQAMAMEQIAmCFCAAAzRAgAYIYIAQDMECEAgBkiBAAwQ4QAAGaIEADADBECAJghQgAAM0QIAGCGCAEAzBAhAIAZIgQAMEOEAABmiBAAwAwRAgCYIUIAADNECABghggBAMwQIQCAGSIEADBDhAAAZogQAMAMEQIAmBlnPQAAeDFunPcvWz/+8Y+Teq2//OUvSa3Dw+NMCABghggBAMwQIQCAGSIEADBDhAAAZogQAMAMEQIAmCFCAAAzRAgAYIYIAQDMECEAgBkiBAAwwwVMAWSUMWO8/935V7/6VVKvxQVM048zIQCAGSIEADDjKUL19fWaO3euAoGA8vPztXLlSl26dClhG+ec6urqVFhYqJycHJWXl+vixYspHRoAMDp4ilBzc7Nqamp09uxZNTY26vbt26qoqFBfX198m+3bt2vnzp1qaGhQa2urQqGQli1bpt7e3pQPDwDIbJ4+mPD+++8n3N+zZ4/y8/N17tw5LVq0SM45vfHGG9qyZYuqqqokSXv37lVBQYEOHDigdevWpW5yAEDGe6T3hCKRiCQpLy9PktTe3q6uri5VVFTEt/H7/Vq8eLHOnDlz398jFospGo0m3AAA2SHpCDnnVFtbq4ULF6q0tFSS1NXVJUkqKChI2LagoCD+3FfV19crGAzGb0VFRcmOBADIMElHaP369frkk0/05z//edBzPp8v4b5zbtBj92zevFmRSCR+6+joSHYkAECGSeqHVTds2KAjR46opaVFU6dOjT8eCoUkDZwRhcPh+OPd3d2Dzo7u8fv98vv9yYwBAMhwns6EnHNav369Dh06pJMnT6qkpCTh+ZKSEoVCITU2NsYf6+/vV3Nzs8rKylIzMQBg1PB0JlRTU6MDBw7ovffeUyAQiL/PEwwGlZOTI5/Pp40bN2rbtm2aPn26pk+frm3btmnSpEl6/vnn0/IHAABkLk8R2rVrlySpvLw84fE9e/ZozZo1kqRNmzbp5s2bevnll/X5559r3rx5OnHihAKBQEoGBgCMHj7nnLMe4v9Fo1EFg0HrMYCs9tZbb3le88ILL6RhElvJXCwV/xOJRJSbmzvkNuxhAIAZIgQAMEOEAABmiBAAwAwRAgCYIUIAADNECABghggBAMwQIQCAGSIEADBDhAAAZogQAMAMEQIAmCFCAAAzRAgAYIYIAQDMECEAgBkiBAAwQ4QAAGaIEADADBECAJghQgAAM0QIAGCGCAEAzBAhAIAZIgQAMDPOegAAI8+NGzc8r7l165bnNRMnTvS8Jhl9fX3D8jrwjjMhAIAZIgQAMEOEAABmiBAAwAwRAgCYIUIAADNECABghggBAMwQIQCAGSIEADBDhAAAZogQAMAMFzAFMMgrr7zieY1zzvOan//8557XJHMx0u985zue12B4cCYEADBDhAAAZogQAMAMEQIAmCFCAAAzRAgAYIYIAQDMECEAgBkiBAAwQ4QAAGaIEADADBECAJjhAqYAUmLz5s2e1/z73//2vKa5udnzmr///e+e12B4cCYEADBDhAAAZjxFqL6+XnPnzlUgEFB+fr5WrlypS5cuJWyzZs0a+Xy+hNv8+fNTOjQAYHTwFKHm5mbV1NTo7Nmzamxs1O3bt1VRUTHoH5lavny5Ojs747djx46ldGgAwOjg6YMJ77//fsL9PXv2KD8/X+fOndOiRYvij/v9foVCodRMCAAYtR7pPaFIJCJJysvLS3i8qalJ+fn5mjFjhtauXavu7u4H/h6xWEzRaDThBgDIDklHyDmn2tpaLVy4UKWlpfHHKysrtX//fp08eVI7duxQa2urli5dqlgsdt/fp76+XsFgMH4rKipKdiQAQIbxOedcMgtramp09OhRffjhh5o6deoDt+vs7FRxcbHeeecdVVVVDXo+FoslBCoajRIiIAPl5OR4XrN+/XrPa/g5ocwRiUSUm5s75DZJ/bDqhg0bdOTIEbW0tAwZIEkKh8MqLi5WW1vbfZ/3+/3y+/3JjAEAyHCeIuSc04YNG/Tuu++qqalJJSUlX7ump6dHHR0dCofDSQ8JABidPL0nVFNTo7ffflsHDhxQIBBQV1eXurq6dPPmTUnSjRs39Nprr+mvf/2rrly5oqamJq1YsUJTpkzRM888k5Y/AAAgc3k6E9q1a5ckqby8POHxPXv2aM2aNRo7dqwuXLigffv26YsvvlA4HNaSJUt08OBBBQKBlA0NABgdPH87big5OTk6fvz4Iw0EAMgeSX86Ll2i0aiCwaD1GACAR/Qwn47jAqYAADNECABghggBAMwQIQCAGSIEADBDhAAAZogQAMAMEQIAmCFCAAAzRAgAYIYIAQDMECEAgBkiBAAwQ4QAAGaIEADADBECAJghQgAAM0QIAGCGCAEAzBAhAIAZIgQAMEOEAABmiBAAwAwRAgCYGXERcs5ZjwAASIGH+Xo+4iLU29trPQIAIAUe5uu5z42wU4+7d+/q2rVrCgQC8vl8Cc9Fo1EVFRWpo6NDubm5RhPaYz8MYD8MYD8MYD8MGAn7wTmn3t5eFRYWasyYoc91xg3TTA9tzJgxmjp16pDb5ObmZvVBdg/7YQD7YQD7YQD7YYD1fggGgw+13Yj7dhwAIHsQIQCAmYyKkN/v19atW+X3+61HMcV+GMB+GMB+GMB+GJBp+2HEfTABAJA9MupMCAAwuhAhAIAZIgQAMEOEAABmMipCb775pkpKSjRx4kTNnj1bp0+fth5pWNXV1cnn8yXcQqGQ9Vhp19LSohUrVqiwsFA+n0+HDx9OeN45p7q6OhUWFionJ0fl5eW6ePGizbBp9HX7Yc2aNYOOj/nz59sMmyb19fWaO3euAoGA8vPztXLlSl26dClhm2w4Hh5mP2TK8ZAxETp48KA2btyoLVu26Pz583r66adVWVmpq1evWo82rJ588kl1dnbGbxcuXLAeKe36+vo0a9YsNTQ03Pf57du3a+fOnWpoaFBra6tCoZCWLVs26q5D+HX7QZKWL1+ecHwcO3ZsGCdMv+bmZtXU1Ojs2bNqbGzU7du3VVFRob6+vvg22XA8PMx+kDLkeHAZ4tvf/rZ78cUXEx771re+5X7xi18YTTT8tm7d6mbNmmU9hilJ7t13343fv3v3rguFQu7111+PP3br1i0XDAbd73//e4MJh8dX94NzzlVXV7sf/vCHJvNY6e7udpJcc3Ozcy57j4ev7gfnMud4yIgzof7+fp07d04VFRUJj1dUVOjMmTNGU9loa2tTYWGhSkpK9Nxzz+ny5cvWI5lqb29XV1dXwrHh9/u1ePHirDs2JKmpqUn5+fmaMWOG1q5dq+7ubuuR0ioSiUiS8vLyJGXv8fDV/XBPJhwPGRGh69ev686dOyooKEh4vKCgQF1dXUZTDb958+Zp3759On78uHbv3q2uri6VlZWpp6fHejQz9/77Z/uxIUmVlZXav3+/Tp48qR07dqi1tVVLly5VLBazHi0tnHOqra3VwoULVVpaKik7j4f77Qcpc46HEXcV7aF89Z92cM4Nemw0q6ysjP965syZWrBggaZNm6a9e/eqtrbWcDJ72X5sSNKqVavivy4tLdWcOXNUXFyso0ePqqqqynCy9Fi/fr0++eQTffjhh4Oey6bj4UH7IVOOh4w4E5oyZYrGjh076G8y3d3dg/7Gk00mT56smTNnqq2tzXoUM/c+HcixMVg4HFZxcfGoPD42bNigI0eO6NSpUwn/9Eu2HQ8P2g/3M1KPh4yI0IQJEzR79mw1NjYmPN7Y2KiysjKjqezFYjF9+umnCofD1qOYKSkpUSgUSjg2+vv71dzcnNXHhiT19PSoo6NjVB0fzjmtX79ehw4d0smTJ1VSUpLwfLYcD1+3H+5nxB4Phh+K8OSdd95x48ePd2+99Zb7xz/+4TZu3OgmT57srly5Yj3asHn11VddU1OTu3z5sjt79qz7wQ9+4AKBwKjfB729ve78+fPu/PnzTpLbuXOnO3/+vPvnP//pnHPu9ddfd8Fg0B06dMhduHDBrV692oXDYReNRo0nT62h9kNvb6979dVX3ZkzZ1x7e7s7deqUW7BggfvmN785qvbDSy+95ILBoGtqanKdnZ3x25dffhnfJhuOh6/bD5l0PGRMhJxz7ne/+50rLi52EyZMcE899VTCxxGzwapVq1w4HHbjx493hYWFrqqqyl28eNF6rLQ7deqUkzToVl1d7Zwb+Fju1q1bXSgUcn6/3y1atMhduHDBdug0GGo/fPnll66iosI99thjbvz48e7xxx931dXV7urVq9Zjp9T9/vyS3J49e+LbZMPx8HX7IZOOB/4pBwCAmYx4TwgAMDoRIQCAGSIEADBDhAAAZogQAMAMEQIAmCFCAAAzRAgAYIYIAQDMECEAgBkiBAAwQ4QAAGb+C0beoK+lr26zAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# We can check the dataloader\n",
    "_, (example_datas, labels) = next(enumerate(test_loader))\n",
    "sample = example_datas[0][0]\n",
    "# show the data\n",
    "plt.imshow(sample, cmap='gray', interpolation='none')\n",
    "print(\"Label: \"+ str(labels[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Now we can start to build our CNN model\n",
    "\n",
    "## Then define the model class\n",
    "class CNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(CNN, self).__init__()\n",
    "        #input channel 1, output channel 10\n",
    "        self.conv1 = nn.Conv2d(1, 10, kernel_size=5, stride=1)\n",
    "        #input channel 10, output channel 20\n",
    "        self.conv2 = nn.Conv2d(10, 20, kernel_size=5, stride=1)\n",
    "        #dropout layer\n",
    "        self.conv2_drop = nn.Dropout2d()\n",
    "        #fully connected layer\n",
    "        self.fc1 = nn.Linear(320, 50)\n",
    "        self.fc2 = nn.Linear(50, 10)\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = F.relu(x)\n",
    "        x = self.conv2(x)\n",
    "        x = self.conv2_drop(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = F.relu(x)\n",
    "        x = x.view(-1, 320)\n",
    "        x = self.fc1(x)\n",
    "        x = F.relu(x)\n",
    "        x = F.dropout(x)\n",
    "        x = self.fc2(x)\n",
    "        return F.log_softmax(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "## create model and optimizer\n",
    "learning_rate = 0.01\n",
    "momentum = 0.5\n",
    "device = \"cpu\"\n",
    "model = CNN().to(device) #using cpu here\n",
    "optimizer = optim.SGD(model.parameters(), lr=learning_rate,\n",
    "                      momentum=momentum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "##define train function\n",
    "def train(model, device, train_loader, optimizer, epoch, log_interval=10000):\n",
    "    model.train()\n",
    "    tk0 = tqdm(train_loader, total=int(len(train_loader)))\n",
    "    counter = 0\n",
    "    for batch_idx, (data, target) in enumerate(tk0):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = F.nll_loss(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        counter += 1\n",
    "        tk0.set_postfix(loss=(loss.item()*data.size(0) / (counter * train_loader.batch_size)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "##define test function\n",
    "def test(model, device, test_loader):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            output = model(data)\n",
    "            test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss\n",
    "            pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability\n",
    "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "\n",
    "    print('\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n'.format(\n",
    "        test_loss, correct, len(test_loader.dataset),\n",
    "        100. * correct / len(test_loader.dataset)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fb4557b43f0a4fd3882297fa87ec7a34",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/938 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/tr/3zw9td7531z0_hhwc8105rv00000gn/T/ipykernel_85088/2730739436.py:29: UserWarning: Implicit dimension choice for log_softmax has been deprecated. Change the call to include dim=X as an argument.\n",
      "  return F.log_softmax(x)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Test set: Average loss: 0.3175, Accuracy: 9028/10000 (90%)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "37dc936b4ce14c369bde173daa5ffc90",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/938 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Test set: Average loss: 0.2164, Accuracy: 9356/10000 (94%)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "453f72ee32fb4ba4acf66c5cee078bd5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/938 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Test set: Average loss: 0.1732, Accuracy: 9491/10000 (95%)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "num_epoch = 3\n",
    "for epoch in range(1, num_epoch + 1):\n",
    "        train(model, device, train_loader, optimizer, epoch)\n",
    "        test(model, device, test_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1           [-1, 10, 24, 24]             260\n",
      "            Conv2d-2             [-1, 20, 8, 8]           5,020\n",
      "         Dropout2d-3             [-1, 20, 8, 8]               0\n",
      "            Linear-4                   [-1, 50]          16,050\n",
      "            Linear-5                   [-1, 10]             510\n",
      "================================================================\n",
      "Total params: 21,840\n",
      "Trainable params: 21,840\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.00\n",
      "Forward/backward pass size (MB): 0.06\n",
      "Params size (MB): 0.08\n",
      "Estimated Total Size (MB): 0.15\n",
      "----------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/tr/3zw9td7531z0_hhwc8105rv00000gn/T/ipykernel_85088/2730739436.py:29: UserWarning: Implicit dimension choice for log_softmax has been deprecated. Change the call to include dim=X as an argument.\n",
      "  return F.log_softmax(x)\n"
     ]
    }
   ],
   "source": [
    "summary(model, (1, 28, 28))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "FiltrationsForGNNs",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}