example_data/model_generation/MNIST/generate_model_binary.ipynb

{
 "metadata": {
  "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.8.5"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.8.5 64-bit ('base': conda)"
  },
  "interpreter": {
   "hash": "90e02b2587fbb2ca467fc85381f5522fddb4a9e5fbb8605712260c849ecf752b"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "source": [
    "## Train model for binary MNIST\n",
    "**Function        : Train model in pytorch for binary MNIST **<br>\n",
    "**Author          : Team DIANNA **<br>\n",
    "**Contributor     : **<br>\n",
    "**First Built     : 2021.06.06 **<br>\n",
    "**Last Update     : 2021.06.26 **<br>\n",
    "**Library         : os, numpy, matplotlib, torch, tensorflow, wandb **<br>\n",
    "**Description     : In this notebook we train models in pytorch for binary MNIST dataset. The trained models will be used to explore the XAI methods later.**<br>\n",
    "**Return Values   : pytorch models and training status (.pt) / training report from weights & biases**<br>\n",
    "**Note**          : The best trained model is stored as ONNX. For more information about how to load an ONNX model in pytorch, visit: https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html <br>"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import os\n",
    "import time as tt\n",
    "import numpy as np\n",
    "# DL framework\n",
    "import torch\n",
    "import torchvision\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torch.utils.data\n",
    "import torch.onnx\n",
    "# for plotting\n",
    "import matplotlib.pyplot as plt\n",
    "# report and monitoring with Weights & Biases\n",
    "import wandb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#################################################################################\n",
    "#########                     path to the dataset                        ########\n",
    "#################################################################################\n",
    "# please specify data path\n",
    "datapath = '/mnt/d/NLeSC/DIANNA/data/mnist/binary-MNIST'\n",
    "# please specify output path\n",
    "output_path = '/mnt/d/NLeSC/DIANNA/codebase/dianna/example_data/model_generation/MNIST'\n",
    "if not os.path.exists(output_path):\n",
    "    os.makedirs(output_path, exist_ok = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "dimensions of mnist:\ndimensions or training set (12665, 784)\ndimensions or training set label (12665,)\ndimensions or testing set (2115, 784)\ndimensions or testing set label (2115,)\nstatistics of training set:\nDigits: 0 1\nlabels: [0 1]\nClass distribution: [5923 6742]\nLabels of training set [0 1 1 1 1 0 1 1 0 0 1 0 0 1 0 1 0 0 1 1]\n"
     ]
    }
   ],
   "source": [
    "#################################################################################\n",
    "#########              extract binary MNIST dataset locally              ########\n",
    "#################################################################################\n",
    "# load binary MNIST from local\n",
    "# load data\n",
    "fd = np.load(os.path.join(datapath, 'binary-mnist.npz'))\n",
    "# training set\n",
    "train_X = fd['X_train']\n",
    "train_y = fd['y_train']\n",
    "# testing set\n",
    "test_X = fd['X_test']\n",
    "test_y = fd['y_test']\n",
    "fd.close()\n",
    "\n",
    "# dimensions of data\n",
    "print(\"dimensions of mnist:\")\n",
    "print(\"dimensions or training set\", train_X.shape)\n",
    "print(\"dimensions or training set label\", train_y.shape)\n",
    "print(\"dimensions or testing set\", test_X.shape)\n",
    "print(\"dimensions or testing set label\", test_y.shape)\n",
    "# statistics of training set\n",
    "print(\"statistics of training set:\")\n",
    "print(\"Digits: 0 1\")\n",
    "print(\"labels: {}\".format(np.unique(train_y)))\n",
    "print(\"Class distribution: {}\".format(np.bincount(train_y)))\n",
    "print(\"Labels of training set\", train_y[:20])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "*******************  check the version of pytorch  *********************\n",
      "Pytorch version 1.8.1\n",
      "Is CUDA available? False\n",
      "Device to be used for computation: cpu\n",
      "*******************  login weights & biases  *********************\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mgit-yang\u001b[0m (use `wandb login --relogin` to force relogin)\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "metadata": {},
     "execution_count": 4
    }
   ],
   "source": [
    "#################################################################################\n",
    "#########              set up the environment for training               ########\n",
    "#################################################################################\n",
    "print ('*******************  check the version of pytorch  *********************')\n",
    "print (\"Pytorch version {}\".format(torch.__version__))\n",
    "# check if CUDA is available\n",
    "use_cuda = torch.cuda.is_available()\n",
    "print(\"Is CUDA available? {}\".format(use_cuda))\n",
    "# use GPU if possible\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(\"Device to be used for computation: {}\".format(device))\n",
    "print ('*******************  login weights & biases  *********************')\n",
    "# call weights & biases service\n",
    "wandb.login()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#################################################################################\n",
    "#########                      build neural network                      ########\n",
    "#################################################################################\n",
    "class MnistNet(nn.Module):\n",
    "    def __init__(self, kernels=[16, 32], dropout = 0.1, classes=2):\n",
    "        '''\n",
    "        Two layer CNN model with max pooling.\n",
    "        '''\n",
    "        super(MnistNet, self).__init__()\n",
    "        self.kernels = kernels\n",
    "        # 1st layer\n",
    "        self.layer1 = nn.Sequential(\n",
    "            nn.Conv2d(1, kernels[0], kernel_size=5, stride=1, padding=2),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout()\n",
    "        )\n",
    "        # 2nd layer\n",
    "        self.layer2 = nn.Sequential(\n",
    "            nn.Conv2d(kernels[0], kernels[1], kernel_size=5, stride=1, padding=2),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout()\n",
    "        )\n",
    "        self.fc1 = nn.Linear(7 * 7 * kernels[-1], kernels[-1]) # pixel 28 / maxpooling 2 * 2 = 7\n",
    "        self.fc2 = nn.Linear(kernels[-1], classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.layer1(x)\n",
    "        x = self.layer2(x)\n",
    "        x = x.reshape(x.size(0), -1)\n",
    "        x = self.fc1(x)\n",
    "        x = self.fc2(x)\n",
    "\n",
    "        return F.log_softmax(x, dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#################################################################################\n",
    "#########            configure hyper-parameters & prepare data           ########\n",
    "#################################################################################\n",
    "hyperparameters = dict(\n",
    "    epoch = 10,\n",
    "    classes = 2,\n",
    "    kernels = [16, 32],\n",
    "    batch_size = 64,\n",
    "    learning_rate = 0.001,\n",
    "    dropout = 0.5,\n",
    "    dataset = 'MNIST',\n",
    "    architecture = 'CNN'\n",
    ")\n",
    "\n",
    "# initialize weights & biases service\n",
    "#mode = 'online'\n",
    "mode = 'disabled'\n",
    "wandb.init(config=hyperparameters, project='mnist', entity='dianna-ai', mode=mode)\n",
    "config = wandb.config\n",
    "\n",
    "# use pytorch data loader\n",
    "train_X_torch = torch.from_numpy(train_X).type(torch.FloatTensor)\n",
    "train_y_torch = torch.from_numpy(train_y).type(torch.LongTensor)\n",
    "\n",
    "test_X_torch = torch.from_numpy(test_X).type(torch.FloatTensor)\n",
    "test_y_torch = torch.from_numpy(test_y).type(torch.LongTensor)\n",
    "\n",
    "# reshape the input following the definition in pytorch (batch, channel, Height, Width)\n",
    "train_X_torch = train_X_torch.view(-1,1,28,28)\n",
    "test_X_torch = test_X_torch.view(-1,1,28,28)\n",
    "\n",
    "# pytorch train and test sets\n",
    "train_set = torch.utils.data.TensorDataset(train_X_torch,train_y_torch)\n",
    "test_set = torch.utils.data.TensorDataset(test_X_torch,test_y_torch)\n",
    "\n",
    "# data loader\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size = config.batch_size, shuffle = False)\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size = config.batch_size, shuffle = False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Model details:\n MnistNet(\n  (layer1): Sequential(\n    (0): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n    (1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n    (2): ReLU()\n    (3): Dropout(p=0.5, inplace=False)\n  )\n  (layer2): Sequential(\n    (0): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n    (1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n    (2): ReLU()\n    (3): Dropout(p=0.5, inplace=False)\n  )\n  (fc1): Linear(in_features=1568, out_features=32, bias=True)\n  (fc2): Linear(in_features=32, out_features=2, bias=True)\n)\nOptimizer details:\n Adam (\nParameter Group 0\n    amsgrad: False\n    betas: (0.9, 0.999)\n    eps: 1e-08\n    lr: 0.001\n    weight_decay: 0\n)\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "[<wandb.wandb_torch.TorchGraph at 0x7fcda9297430>]"
      ]
     },
     "metadata": {},
     "execution_count": 8
    }
   ],
   "source": [
    "#################################################################################\n",
    "#########               create model and choose optimizer                ########\n",
    "#################################################################################\n",
    "model = MnistNet(config.kernels, config.dropout, config.classes).to(device)\n",
    "criterion = nn.CrossEntropyLoss().to(device)\n",
    "#criterion = nn.BCELoss()\n",
    "#criterion = nn.BCEWithLogitsLoss() # same as BCELoss(torch.sigmoid(x),...), but more numerically stable\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)\n",
    "print('Model details:\\n', model)\n",
    "print('Optimizer details:\\n',optimizer)\n",
    "\n",
    "wandb.watch(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Epoch : 0 [0/12665 (0%)]\tLoss: 8.398658\t Accuracy:51.562%\n",
      "Epoch : 0 [1280/12665 (10%)]\tLoss: 0.054485\t Accuracy:92.783%\n",
      "Epoch : 0 [2560/12665 (20%)]\tLoss: 0.000000\t Accuracy:95.770%\n",
      "Epoch : 0 [3840/12665 (30%)]\tLoss: 0.000000\t Accuracy:96.952%\n",
      "Epoch : 0 [5120/12665 (40%)]\tLoss: 0.000000\t Accuracy:97.589%\n",
      "Epoch : 0 [6400/12665 (51%)]\tLoss: 0.089167\t Accuracy:97.912%\n",
      "Epoch : 0 [7680/12665 (61%)]\tLoss: 0.000000\t Accuracy:98.166%\n",
      "Epoch : 0 [8960/12665 (71%)]\tLoss: 0.000000\t Accuracy:98.349%\n",
      "Epoch : 0 [10240/12665 (81%)]\tLoss: 0.013582\t Accuracy:98.535%\n",
      "Epoch : 0 [11520/12665 (91%)]\tLoss: 0.000000\t Accuracy:98.679%\n",
      "Epoch : 1 [0/12665 (0%)]\tLoss: 0.380646\t Accuracy:98.438%\n",
      "Epoch : 1 [1280/12665 (10%)]\tLoss: 0.000000\t Accuracy:99.479%\n",
      "Epoch : 1 [2560/12665 (20%)]\tLoss: 0.000000\t Accuracy:99.543%\n",
      "Epoch : 1 [3840/12665 (30%)]\tLoss: 0.008374\t Accuracy:99.616%\n",
      "Epoch : 1 [5120/12665 (40%)]\tLoss: 0.000000\t Accuracy:99.614%\n",
      "Epoch : 1 [6400/12665 (51%)]\tLoss: 0.000000\t Accuracy:99.613%\n",
      "Epoch : 1 [7680/12665 (61%)]\tLoss: 0.000000\t Accuracy:99.626%\n",
      "Epoch : 1 [8960/12665 (71%)]\tLoss: 0.000036\t Accuracy:99.590%\n",
      "Epoch : 1 [10240/12665 (81%)]\tLoss: 0.000000\t Accuracy:99.612%\n",
      "Epoch : 1 [11520/12665 (91%)]\tLoss: 0.000000\t Accuracy:99.646%\n",
      "Epoch : 2 [0/12665 (0%)]\tLoss: 0.003869\t Accuracy:100.000%\n",
      "Epoch : 2 [1280/12665 (10%)]\tLoss: 0.000000\t Accuracy:99.926%\n",
      "Epoch : 2 [2560/12665 (20%)]\tLoss: 0.000000\t Accuracy:99.924%\n",
      "Epoch : 2 [3840/12665 (30%)]\tLoss: 0.000000\t Accuracy:99.898%\n",
      "Epoch : 2 [5120/12665 (40%)]\tLoss: 0.000000\t Accuracy:99.884%\n",
      "Epoch : 2 [6400/12665 (51%)]\tLoss: 0.000000\t Accuracy:99.861%\n",
      "Epoch : 2 [7680/12665 (61%)]\tLoss: 0.000000\t Accuracy:99.832%\n",
      "Epoch : 2 [8960/12665 (71%)]\tLoss: 0.000000\t Accuracy:99.801%\n",
      "Epoch : 2 [10240/12665 (81%)]\tLoss: 0.000000\t Accuracy:99.728%\n",
      "Epoch : 2 [11520/12665 (91%)]\tLoss: 0.000000\t Accuracy:99.724%\n",
      "Epoch : 3 [0/12665 (0%)]\tLoss: 0.000000\t Accuracy:100.000%\n",
      "Epoch : 3 [1280/12665 (10%)]\tLoss: 0.006490\t Accuracy:99.851%\n",
      "Epoch : 3 [2560/12665 (20%)]\tLoss: 0.000000\t Accuracy:99.809%\n",
      "Epoch : 3 [3840/12665 (30%)]\tLoss: 0.000013\t Accuracy:99.769%\n",
      "Epoch : 3 [5120/12665 (40%)]\tLoss: 0.000000\t Accuracy:99.769%\n",
      "Epoch : 3 [6400/12665 (51%)]\tLoss: 0.000686\t Accuracy:99.768%\n",
      "Epoch : 3 [7680/12665 (61%)]\tLoss: 0.004232\t Accuracy:99.768%\n",
      "Epoch : 3 [8960/12665 (71%)]\tLoss: 0.000011\t Accuracy:99.778%\n",
      "Epoch : 3 [10240/12665 (81%)]\tLoss: 0.005645\t Accuracy:99.806%\n",
      "Epoch : 3 [11520/12665 (91%)]\tLoss: 0.000000\t Accuracy:99.793%\n",
      "Epoch : 4 [0/12665 (0%)]\tLoss: 0.000001\t Accuracy:100.000%\n",
      "Epoch : 4 [1280/12665 (10%)]\tLoss: 0.000000\t Accuracy:99.702%\n",
      "Epoch : 4 [2560/12665 (20%)]\tLoss: 0.000000\t Accuracy:99.695%\n",
      "Epoch : 4 [3840/12665 (30%)]\tLoss: 0.000000\t Accuracy:99.795%\n",
      "Epoch : 4 [5120/12665 (40%)]\tLoss: 0.000000\t Accuracy:99.826%\n",
      "Epoch : 4 [6400/12665 (51%)]\tLoss: 0.024788\t Accuracy:99.830%\n",
      "Epoch : 4 [7680/12665 (61%)]\tLoss: 0.000000\t Accuracy:99.819%\n",
      "Epoch : 4 [8960/12665 (71%)]\tLoss: 0.000000\t Accuracy:99.767%\n",
      "Epoch : 4 [10240/12665 (81%)]\tLoss: 0.000000\t Accuracy:99.777%\n",
      "Epoch : 4 [11520/12665 (91%)]\tLoss: 0.000000\t Accuracy:99.793%\n",
      "Epoch : 5 [0/12665 (0%)]\tLoss: 0.000398\t Accuracy:100.000%\n",
      "Epoch : 5 [1280/12665 (10%)]\tLoss: 0.000027\t Accuracy:100.000%\n",
      "Epoch : 5 [2560/12665 (20%)]\tLoss: 0.000000\t Accuracy:99.962%\n",
      "Epoch : 5 [3840/12665 (30%)]\tLoss: 0.000141\t Accuracy:99.898%\n",
      "Epoch : 5 [5120/12665 (40%)]\tLoss: 0.000000\t Accuracy:99.904%\n",
      "Epoch : 5 [6400/12665 (51%)]\tLoss: 0.000024\t Accuracy:99.892%\n",
      "Epoch : 5 [7680/12665 (61%)]\tLoss: 0.000000\t Accuracy:99.897%\n",
      "Epoch : 5 [8960/12665 (71%)]\tLoss: 0.000002\t Accuracy:99.878%\n",
      "Epoch : 5 [10240/12665 (81%)]\tLoss: 0.000000\t Accuracy:99.893%\n",
      "Epoch : 5 [11520/12665 (91%)]\tLoss: 0.000000\t Accuracy:99.896%\n",
      "Epoch : 6 [0/12665 (0%)]\tLoss: 0.000001\t Accuracy:100.000%\n",
      "Epoch : 6 [1280/12665 (10%)]\tLoss: 0.065525\t Accuracy:99.628%\n",
      "Epoch : 6 [2560/12665 (20%)]\tLoss: 0.000002\t Accuracy:99.543%\n",
      "Epoch : 6 [3840/12665 (30%)]\tLoss: 0.000000\t Accuracy:99.616%\n",
      "Epoch : 6 [5120/12665 (40%)]\tLoss: 0.000000\t Accuracy:99.672%\n",
      "Epoch : 6 [6400/12665 (51%)]\tLoss: 0.017646\t Accuracy:99.706%\n",
      "Epoch : 6 [7680/12665 (61%)]\tLoss: 0.000000\t Accuracy:99.755%\n",
      "Epoch : 6 [8960/12665 (71%)]\tLoss: 0.000000\t Accuracy:99.778%\n",
      "Epoch : 6 [10240/12665 (81%)]\tLoss: 0.000000\t Accuracy:99.806%\n",
      "Epoch : 6 [11520/12665 (91%)]\tLoss: 0.000000\t Accuracy:99.810%\n",
      "Epoch : 7 [0/12665 (0%)]\tLoss: 0.001043\t Accuracy:100.000%\n",
      "Epoch : 7 [1280/12665 (10%)]\tLoss: 0.000009\t Accuracy:99.777%\n",
      "Epoch : 7 [2560/12665 (20%)]\tLoss: 0.000015\t Accuracy:99.886%\n",
      "Epoch : 7 [3840/12665 (30%)]\tLoss: 0.003390\t Accuracy:99.923%\n",
      "Epoch : 7 [5120/12665 (40%)]\tLoss: 0.000000\t Accuracy:99.904%\n",
      "Epoch : 7 [6400/12665 (51%)]\tLoss: 0.001600\t Accuracy:99.892%\n",
      "Epoch : 7 [7680/12665 (61%)]\tLoss: 0.000002\t Accuracy:99.884%\n",
      "Epoch : 7 [8960/12665 (71%)]\tLoss: 0.006719\t Accuracy:99.856%\n",
      "Epoch : 7 [10240/12665 (81%)]\tLoss: 0.000083\t Accuracy:99.864%\n",
      "Epoch : 7 [11520/12665 (91%)]\tLoss: 0.000000\t Accuracy:99.879%\n",
      "Epoch : 8 [0/12665 (0%)]\tLoss: 0.000047\t Accuracy:100.000%\n",
      "Epoch : 8 [1280/12665 (10%)]\tLoss: 0.000000\t Accuracy:100.000%\n",
      "Epoch : 8 [2560/12665 (20%)]\tLoss: 0.000000\t Accuracy:99.924%\n",
      "Epoch : 8 [3840/12665 (30%)]\tLoss: 0.000035\t Accuracy:99.923%\n",
      "Epoch : 8 [5120/12665 (40%)]\tLoss: 0.000000\t Accuracy:99.923%\n",
      "Epoch : 8 [6400/12665 (51%)]\tLoss: 0.111008\t Accuracy:99.892%\n",
      "Epoch : 8 [7680/12665 (61%)]\tLoss: 0.000000\t Accuracy:99.858%\n",
      "Epoch : 8 [8960/12665 (71%)]\tLoss: 0.000393\t Accuracy:99.867%\n",
      "Epoch : 8 [10240/12665 (81%)]\tLoss: 0.000000\t Accuracy:99.874%\n",
      "Epoch : 8 [11520/12665 (91%)]\tLoss: 0.000000\t Accuracy:99.879%\n",
      "Epoch : 9 [0/12665 (0%)]\tLoss: 0.000000\t Accuracy:100.000%\n",
      "Epoch : 9 [1280/12665 (10%)]\tLoss: 0.000000\t Accuracy:99.926%\n",
      "Epoch : 9 [2560/12665 (20%)]\tLoss: 0.000000\t Accuracy:99.962%\n",
      "Epoch : 9 [3840/12665 (30%)]\tLoss: 0.000000\t Accuracy:99.949%\n",
      "Epoch : 9 [5120/12665 (40%)]\tLoss: 0.000000\t Accuracy:99.961%\n",
      "Epoch : 9 [6400/12665 (51%)]\tLoss: 0.000045\t Accuracy:99.954%\n",
      "Epoch : 9 [7680/12665 (61%)]\tLoss: 0.000805\t Accuracy:99.948%\n",
      "Epoch : 9 [8960/12665 (71%)]\tLoss: 0.000001\t Accuracy:99.934%\n",
      "Epoch : 9 [10240/12665 (81%)]\tLoss: 0.000145\t Accuracy:99.942%\n",
      "Epoch : 9 [11520/12665 (91%)]\tLoss: 0.000000\t Accuracy:99.948%\n",
      "The checkpoint of the model and training status is saved.\n",
      "--- 0.686580232779185 minutes ---\n"
     ]
    }
   ],
   "source": [
    "#################################################################################\n",
    "#########                   start the training process                   ########\n",
    "#################################################################################\n",
    "# calculate the time for the code execution\n",
    "start_time = tt.time()\n",
    "# switch model into train mode\n",
    "model.train()\n",
    "for epoch in range(config.epoch):\n",
    "    # number of correction predictions\n",
    "    correct = 0\n",
    "    for batch_idx, (X_batch, y_batch) in enumerate(train_loader):\n",
    "        var_X_batch = torch.autograd.Variable(X_batch).to(device)\n",
    "        var_y_batch = torch.autograd.Variable(y_batch).to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(var_X_batch).squeeze(1)\n",
    "        loss = criterion(output, var_y_batch)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        # Total correct predictions\n",
    "        predicted = torch.max(output.data, 1)[1]\n",
    "        #predicted = torch.round(torch.sigmoid(output))\n",
    "        correct += (predicted == var_y_batch).sum()\n",
    "        # log the training loss and accuracy in wandb\n",
    "        wandb.log({'train_loss': loss.item(), 'train_acc': float(correct*100) / float(config.batch_size*(batch_idx+1))})\n",
    "        if batch_idx % 20 == 0:\n",
    "            print('Epoch : {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}\\t Accuracy:{:.3f}%'.format(\n",
    "                  epoch, batch_idx*len(X_batch), len(train_loader.dataset), 100.* batch_idx / len(train_loader),\n",
    "                  loss.item(), float(correct*100) / float(config.batch_size*(batch_idx+1))))\n",
    "\n",
    "# save the general checkpoint\n",
    "torch.save({\n",
    "            'epoch': epoch,\n",
    "            'model_state_dict': model.state_dict(),\n",
    "            'optimizer_state_dict': optimizer.state_dict(),\n",
    "            'loss': loss.item()\n",
    "            }, os.path.join(output_path,'mnistnet_training_checkpoint.pt'))\n",
    "print(\"The checkpoint of the model and training status is saved.\")\n",
    "\n",
    "print (\"--- %s minutes ---\" % ((tt.time() - start_time)/60))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Test accuracy:100.000% \nTest accuracy:100.000% \nTest accuracy:100.000% \nTest accuracy:100.000% \nTest accuracy:100.000% \nTest accuracy:99.940% \nTest accuracy:99.899% \n"
     ]
    }
   ],
   "source": [
    "#################################################################################\n",
    "#########                  start the evaluation process                  ########\n",
    "#################################################################################\n",
    "# switch model into evaluation mode\n",
    "model.eval()\n",
    "correct = 0\n",
    "for batch_idx, (test_imgs, test_labels) in enumerate(test_loader):\n",
    "    test_imgs = torch.autograd.Variable(test_imgs).float()\n",
    "    output = model(test_imgs)\n",
    "    predicted = torch.max(output,1)[1]\n",
    "    correct += (predicted == test_labels).sum()\n",
    "    wandb.log({'test_acc': float(correct*100) / float(config.batch_size*(batch_idx+1))})\n",
    "    if batch_idx % 5 == 0:\n",
    "        print(\"Test accuracy:{:.3f}% \".format(float(correct*100) / float(config.batch_size*(batch_idx+1))))\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Predictions: [1 0 1]\nGround Truth: tensor([1, 0, 1])\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "<Figure size 288x288 with 3 Axes>",
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<!-- Created with matplotlib (https://matplotlib.org/) -->\n<svg height=\"108.07875pt\" version=\"1.1\" viewBox=\"0 0 280.8 108.07875\" width=\"280.8pt\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n <defs>\n  <style type=\"text/css\">\n*{stroke-linecap:butt;stroke-linejoin:round;}\n  </style>\n </defs>\n <g id=\"figure_1\">\n  <g id=\"patch_1\">\n   <path d=\"M 0 108.07875 \nL 280.8 108.07875 \nL 280.8 0 \nL 0 0 \nz\n\" style=\"fill:none;\"/>\n  </g>\n  <g id=\"axes_1\">\n   <g clip-path=\"url(#pa3021f6c68)\">\n    <image height=\"82\" id=\"image8552007a58\" transform=\"scale(1 -1)translate(0 -82)\" width=\"82\" x=\"7.2\" xlink:href=\"data:image/png;base64,\niVBORw0KGgoAAAANSUhEUgAAAFIAAABSCAYAAADHLIObAAAABHNCSVQICAgIfAhkiAAAAWtJREFUeJzt2rGNg0AQQNHlYiiCFFEAERIRDbgaWqEJt+A2aGYd3YnsfL4v7yz6TyIffY3lXZsmpZST/u2r9ABXYUiIISGGhBgSYkhImJDruqac86/Pvu9pWZY0jmNq27b02D+aFOQcmfN7YzRNA0/ynjAbWTtDQgwJMSTEkBBDQgwJMSTEkBBDQqoP+X0Hn+e56BzV37XPSt67q9/IKAwJMSTEkBBDQgwJMSTEkBBDQgwJMSTEkJAwIfu+T9u2pfv9XnqUt4T59eesxrcuwmxk7QwJMSTEkBBDQgwJMSTkUiFLniMvdSA/+3TUS21kSYaEGBJiSIghIYaEGBJiSIghIYaEGBJiSIghIYaEGBJiSIghIYaEhAx5u91Kj/BnIf+zOavlhaqQG1kjQ0IMCQkf8vF4lB7hZTny03VdHoYhT9OUX3Ucx8fnDP+tXYvwH+1aGBJiSIghIYaEGBJiSIghIYaEGBJiSIghIU9YWLbNEZg6RgAAAABJRU5ErkJggg==\" y=\"-18.87875\"/>\n   </g>\n   <g id=\"text_1\">\n    <!-- tensor(1) -->\n    <defs>\n     <path d=\"M 18.3125 70.21875 \nL 18.3125 54.6875 \nL 36.8125 54.6875 \nL 36.8125 47.703125 \nL 18.3125 47.703125 \nL 18.3125 18.015625 \nQ 18.3125 11.328125 20.140625 9.421875 \nQ 21.96875 7.515625 27.59375 7.515625 \nL 36.8125 7.515625 \nL 36.8125 0 \nL 27.59375 0 \nQ 17.1875 0 13.234375 3.875 \nQ 9.28125 7.765625 9.28125 18.015625 \nL 9.28125 47.703125 \nL 2.6875 47.703125 \nL 2.6875 54.6875 \nL 9.28125 54.6875 \nL 9.28125 70.21875 \nz\n\" id=\"DejaVuSans-116\"/>\n     <path d=\"M 56.203125 29.59375 \nL 56.203125 25.203125 \nL 14.890625 25.203125 \nQ 15.484375 15.921875 20.484375 11.0625 \nQ 25.484375 6.203125 34.421875 6.203125 \nQ 39.59375 6.203125 44.453125 7.46875 \nQ 49.3125 8.734375 54.109375 11.28125 \nL 54.109375 2.78125 \nQ 49.265625 0.734375 44.1875 -0.34375 \nQ 39.109375 -1.421875 33.890625 -1.421875 \nQ 20.796875 -1.421875 13.15625 6.1875 \nQ 5.515625 13.8125 5.515625 26.8125 \nQ 5.515625 40.234375 12.765625 48.109375 \nQ 20.015625 56 32.328125 56 \nQ 43.359375 56 49.78125 48.890625 \nQ 56.203125 41.796875 56.203125 29.59375 \nz\nM 47.21875 32.234375 \nQ 47.125 39.59375 43.09375 43.984375 \nQ 39.0625 48.390625 32.421875 48.390625 \nQ 24.90625 48.390625 20.390625 44.140625 \nQ 15.875 39.890625 15.1875 32.171875 \nz\n\" id=\"DejaVuSans-101\"/>\n     <path d=\"M 54.890625 33.015625 \nL 54.890625 0 \nL 45.90625 0 \nL 45.90625 32.71875 \nQ 45.90625 40.484375 42.875 44.328125 \nQ 39.84375 48.1875 33.796875 48.1875 \nQ 26.515625 48.1875 22.3125 43.546875 \nQ 18.109375 38.921875 18.109375 30.90625 \nL 18.109375 0 \nL 9.078125 0 \nL 9.078125 54.6875 \nL 18.109375 54.6875 \nL 18.109375 46.1875 \nQ 21.34375 51.125 25.703125 53.5625 \nQ 30.078125 56 35.796875 56 \nQ 45.21875 56 50.046875 50.171875 \nQ 54.890625 44.34375 54.890625 33.015625 \nz\n\" id=\"DejaVuSans-110\"/>\n     <path d=\"M 44.28125 53.078125 \nL 44.28125 44.578125 \nQ 40.484375 46.53125 36.375 47.5 \nQ 32.28125 48.484375 27.875 48.484375 \nQ 21.1875 48.484375 17.84375 46.4375 \nQ 14.5 44.390625 14.5 40.28125 \nQ 14.5 37.15625 16.890625 35.375 \nQ 19.28125 33.59375 26.515625 31.984375 \nL 29.59375 31.296875 \nQ 39.15625 29.25 43.1875 25.515625 \nQ 47.21875 21.78125 47.21875 15.09375 \nQ 47.21875 7.46875 41.1875 3.015625 \nQ 35.15625 -1.421875 24.609375 -1.421875 \nQ 20.21875 -1.421875 15.453125 -0.5625 \nQ 10.6875 0.296875 5.421875 2 \nL 5.421875 11.28125 \nQ 10.40625 8.6875 15.234375 7.390625 \nQ 20.0625 6.109375 24.8125 6.109375 \nQ 31.15625 6.109375 34.5625 8.28125 \nQ 37.984375 10.453125 37.984375 14.40625 \nQ 37.984375 18.0625 35.515625 20.015625 \nQ 33.0625 21.96875 24.703125 23.78125 \nL 21.578125 24.515625 \nQ 13.234375 26.265625 9.515625 29.90625 \nQ 5.8125 33.546875 5.8125 39.890625 \nQ 5.8125 47.609375 11.28125 51.796875 \nQ 16.75 56 26.8125 56 \nQ 31.78125 56 36.171875 55.265625 \nQ 40.578125 54.546875 44.28125 53.078125 \nz\n\" id=\"DejaVuSans-115\"/>\n     <path d=\"M 30.609375 48.390625 \nQ 23.390625 48.390625 19.1875 42.75 \nQ 14.984375 37.109375 14.984375 27.296875 \nQ 14.984375 17.484375 19.15625 11.84375 \nQ 23.34375 6.203125 30.609375 6.203125 \nQ 37.796875 6.203125 41.984375 11.859375 \nQ 46.1875 17.53125 46.1875 27.296875 \nQ 46.1875 37.015625 41.984375 42.703125 \nQ 37.796875 48.390625 30.609375 48.390625 \nz\nM 30.609375 56 \nQ 42.328125 56 49.015625 48.375 \nQ 55.71875 40.765625 55.71875 27.296875 \nQ 55.71875 13.875 49.015625 6.21875 \nQ 42.328125 -1.421875 30.609375 -1.421875 \nQ 18.84375 -1.421875 12.171875 6.21875 \nQ 5.515625 13.875 5.515625 27.296875 \nQ 5.515625 40.765625 12.171875 48.375 \nQ 18.84375 56 30.609375 56 \nz\n\" id=\"DejaVuSans-111\"/>\n     <path d=\"M 41.109375 46.296875 \nQ 39.59375 47.171875 37.8125 47.578125 \nQ 36.03125 48 33.890625 48 \nQ 26.265625 48 22.1875 43.046875 \nQ 18.109375 38.09375 18.109375 28.8125 \nL 18.109375 0 \nL 9.078125 0 \nL 9.078125 54.6875 \nL 18.109375 54.6875 \nL 18.109375 46.1875 \nQ 20.953125 51.171875 25.484375 53.578125 \nQ 30.03125 56 36.53125 56 \nQ 37.453125 56 38.578125 55.875 \nQ 39.703125 55.765625 41.0625 55.515625 \nz\n\" id=\"DejaVuSans-114\"/>\n     <path d=\"M 31 75.875 \nQ 24.46875 64.65625 21.28125 53.65625 \nQ 18.109375 42.671875 18.109375 31.390625 \nQ 18.109375 20.125 21.3125 9.0625 \nQ 24.515625 -2 31 -13.1875 \nL 23.1875 -13.1875 \nQ 15.875 -1.703125 12.234375 9.375 \nQ 8.59375 20.453125 8.59375 31.390625 \nQ 8.59375 42.28125 12.203125 53.3125 \nQ 15.828125 64.359375 23.1875 75.875 \nz\n\" id=\"DejaVuSans-40\"/>\n     <path d=\"M 12.40625 8.296875 \nL 28.515625 8.296875 \nL 28.515625 63.921875 \nL 10.984375 60.40625 \nL 10.984375 69.390625 \nL 28.421875 72.90625 \nL 38.28125 72.90625 \nL 38.28125 8.296875 \nL 54.390625 8.296875 \nL 54.390625 0 \nL 12.40625 0 \nz\n\" id=\"DejaVuSans-49\"/>\n     <path d=\"M 8.015625 75.875 \nL 15.828125 75.875 \nQ 23.140625 64.359375 26.78125 53.3125 \nQ 30.421875 42.28125 30.421875 31.390625 \nQ 30.421875 20.453125 26.78125 9.375 \nQ 23.140625 -1.703125 15.828125 -13.1875 \nL 8.015625 -13.1875 \nQ 14.5 -2 17.703125 9.0625 \nQ 20.90625 20.125 20.90625 31.390625 \nQ 20.90625 42.671875 17.703125 53.65625 \nQ 14.5 64.65625 8.015625 75.875 \nz\n\" id=\"DejaVuSans-41\"/>\n    </defs>\n    <g transform=\"translate(29.59375 13.27875)scale(0.08 -0.08)\">\n     <use xlink:href=\"#DejaVuSans-116\"/>\n     <use x=\"39.208984\" xlink:href=\"#DejaVuSans-101\"/>\n     <use x=\"100.732422\" xlink:href=\"#DejaVuSans-110\"/>\n     <use x=\"164.111328\" xlink:href=\"#DejaVuSans-115\"/>\n     <use x=\"216.210938\" xlink:href=\"#DejaVuSans-111\"/>\n     <use x=\"277.392578\" xlink:href=\"#DejaVuSans-114\"/>\n     <use x=\"318.505859\" xlink:href=\"#DejaVuSans-40\"/>\n     <use x=\"357.519531\" xlink:href=\"#DejaVuSans-49\"/>\n     <use x=\"421.142578\" xlink:href=\"#DejaVuSans-41\"/>\n    </g>\n   </g>\n  </g>\n  <g id=\"axes_2\">\n   <g clip-path=\"url(#pd96658de3c)\">\n    <image height=\"82\" id=\"imagee6b0f586aa\" transform=\"scale(1 -1)translate(0 -82)\" width=\"82\" x=\"99.6\" xlink:href=\"data:image/png;base64,\niVBORw0KGgoAAAANSUhEUgAAAFIAAABSCAYAAADHLIObAAAABHNCSVQICAgIfAhkiAAABABJREFUeJztnLFL+0AUx1/aQFECWmkHRVFQlCJdShEKdejkquDaQZwE/wxHQUEQFBRcxE3QqTgIiuDURaGDi0jRpdaWZiiovd/gr1Jsm15yL3eX5j5wS3J37/XLe3m55FINAAgomAmIdqBfUEIioYREQgmJhBISCSUkEkpIJJSQSCghkdCZJ9B1mJmZgVgsBtFoFLLZLKTTaQzf2vj4+IC9vT14eHiAYrEIhUIBqtWqK7acQFja2toaEQmr/1hNA8a19vPzM0xOTrJMwYSmacJst+IotYPBIMRiMUgkEkJFBAD4CcofKpUKXF5ewuHhIRSLRajValCpVOD7+5uPL3ZbIpEQms52SCaTXFLbUdVOpVJOhgnBrcLXCdvqX1xciA40aqrVqrzFhhDbQ4TCoyCpG3IkfCEkIaStzc7OotrwRWp3AzPlfRGRPPC1kM00n5qaYp7L16nd5OXlhXmFpoT8D+v10tep3UozzcPhsKPxSsg/LC8vOxqnhPzD0tKSo3HSXSN1Xe/42EvTNGg0Gq7ZbdJoNCAYDNoeJ11EGobR8TivAhcIOJNEOiEHBga6nhseHoZ4PA6pVAqOjo44etUb6VJbJiKRCLy/v1P1lS4iZSKTyVD3dSRkOByGfD7vZKiniMfj1H2Z3yLKlOatqxMMvwgh1MVHpbYFdpaNnhcynU7D3NxcW7VPJpNc/WDesiKau7u7jsdrtRpXPzwfkd14e3uD6+trbvY8WWwMw4Cvry/4/PykWjay+Eh7nfRkRJqmCfV6nctWlObjtfHxcct+nhRSBL12bHheyFAoBLqug6Zprm4EmJiYsDzv+apdr9c7HscWdWFhoWcfpj0vstLq49nZGfqcf5snqzYN8/PzUC6XoVQq/b5ynZ6ehlwu53hOqyjvWyHdwEpIzxcbWVBCIqGERIJZyI2NDQw/PA9zsWml3wuPKjYcQBXy+PgYczqpKJVKludRU7uVfkvz+/t7y89iVGpTcn5+bnneNSF3d3fdmloIJycnPfu4/jHP2NgY2dnZIeVyGeXhAU/y+TxZXV0lmqbhf0JnF9M04ebmBg4ODniYQ+X09BSenp6o+roekYFAgBiGQSKRCNne3hYdZLYYGRkhoVDInU/osHh9fYXR0VFR5qnwxMuv/f19keZ7YvdvHlxP7W5taGiIZDIZsr6+Tm5vb0VncRubm5t2fo84IWmaSFZWVujrAGXU+pJCoUDdV3ohFxcX4erqCkzT5Gbz8fERstks9W0PgItrbV4QDttRaJA+InuxtbUl2oVfhBcUljY4OGi7iORyORKNRlH98Hxqy4LnU1sWlJBIKCGRUEIioYREQgmJhBISCSUkEkpIJJSQSCghkVBCIvEPbYz0DarzaZ0AAAAASUVORK5CYII=\" y=\"-18.87875\"/>\n   </g>\n   <g id=\"text_2\">\n    <!-- tensor(0) -->\n    <defs>\n     <path d=\"M 31.78125 66.40625 \nQ 24.171875 66.40625 20.328125 58.90625 \nQ 16.5 51.421875 16.5 36.375 \nQ 16.5 21.390625 20.328125 13.890625 \nQ 24.171875 6.390625 31.78125 6.390625 \nQ 39.453125 6.390625 43.28125 13.890625 \nQ 47.125 21.390625 47.125 36.375 \nQ 47.125 51.421875 43.28125 58.90625 \nQ 39.453125 66.40625 31.78125 66.40625 \nz\nM 31.78125 74.21875 \nQ 44.046875 74.21875 50.515625 64.515625 \nQ 56.984375 54.828125 56.984375 36.375 \nQ 56.984375 17.96875 50.515625 8.265625 \nQ 44.046875 -1.421875 31.78125 -1.421875 \nQ 19.53125 -1.421875 13.0625 8.265625 \nQ 6.59375 17.96875 6.59375 36.375 \nQ 6.59375 54.828125 13.0625 64.515625 \nQ 19.53125 74.21875 31.78125 74.21875 \nz\n\" id=\"DejaVuSans-48\"/>\n    </defs>\n    <g transform=\"translate(121.99375 13.27875)scale(0.08 -0.08)\">\n     <use xlink:href=\"#DejaVuSans-116\"/>\n     <use x=\"39.208984\" xlink:href=\"#DejaVuSans-101\"/>\n     <use x=\"100.732422\" xlink:href=\"#DejaVuSans-110\"/>\n     <use x=\"164.111328\" xlink:href=\"#DejaVuSans-115\"/>\n     <use x=\"216.210938\" xlink:href=\"#DejaVuSans-111\"/>\n     <use x=\"277.392578\" xlink:href=\"#DejaVuSans-114\"/>\n     <use x=\"318.505859\" xlink:href=\"#DejaVuSans-40\"/>\n     <use x=\"357.519531\" xlink:href=\"#DejaVuSans-48\"/>\n     <use x=\"421.142578\" xlink:href=\"#DejaVuSans-41\"/>\n    </g>\n   </g>\n  </g>\n  <g id=\"axes_3\">\n   <g clip-path=\"url(#pd8fe64a99a)\">\n    <image height=\"82\" id=\"imageecfe4ecaf6\" transform=\"scale(1 -1)translate(0 -82)\" width=\"82\" x=\"192\" xlink:href=\"data:image/png;base64,\niVBORw0KGgoAAAANSUhEUgAAAFIAAABSCAYAAADHLIObAAAABHNCSVQICAgIfAhkiAAAAodJREFUeJzt2z2K8lAUxvFjCCpaWigIYqmNhaBYWgluIr2FiFuwEezcgoUgli7BStxAsFPsLcRK8UzzBjK8zowfT+5HOD9Ip3Pv/Cfm3mQwQURM4mOO7gnEhYQEkZAgEhJEQoJISBAJCSIhQSQkiLKQjuOQ67rUbreJmf88ZrMZpdNpSiaT5Dh2/L056iOTyXCr1WLP8/gV6/Wap9Mp12q1yOcIOKIfJJ/P82g0Yt/3XwoZ6PV6uiOZEbLZbL4VMHA6nXRH+vNIBDWjxPz5EIlEAjCT6NhxFbeANSH532re7XZ1T+Uhaz7aYSZ+zK05I00nIUEkJIiEBLEy5GQyIc/zqFAo6J7KN5Hv+rfb7Ud3Nj9ZLBba72hCR/SD5HI57vf7vFqt4DENCKjuFjGMY7qntPIaaSIJCSIhQZSHPB6PqodUQvlik81mqVqtUrlcpuVy+fHPM2WxUR4yDLGCmxJSrpEgEhLE+pBsyJNz66+RYTqvl1rPyFQqRcVikSqVis5pQGg9I8NsX8Gtv0aaQkKCSEgQCQkiIUFiFTLYnA8GA+Vjx2r7E6Z6K+QqHe0X4V8cHVWFWH20dZKQIBISREKCSEgQCQkiIUEkJIgxG/IwGzfnckaCSEgQCQkiIUEkJEhsQwYPeev1upLxjHmw+xNb/t8d2zNSNQkJIiFBjA85HA51T+Epxi82gVKpRPv9/q33ymITcr1edU/hV9aEvN1ub72v0WiAZ/KYkY/RHrnf7y+9frfb0Xw+p8PhENGMvrMm5OVyeen14/GYfN+n8/kc0Yz+p/0rus8em83mqa8edzoddl2XHceJ79eM48qaxcZ0EhJEQoJISBAJCSIhQSQkiIQEkZAgX/RBmvr+kTshAAAAAElFTkSuQmCC\" y=\"-18.87875\"/>\n   </g>\n   <g id=\"text_3\">\n    <!-- tensor(1) -->\n    <g transform=\"translate(214.39375 13.27875)scale(0.08 -0.08)\">\n     <use xlink:href=\"#DejaVuSans-116\"/>\n     <use x=\"39.208984\" xlink:href=\"#DejaVuSans-101\"/>\n     <use x=\"100.732422\" xlink:href=\"#DejaVuSans-110\"/>\n     <use x=\"164.111328\" xlink:href=\"#DejaVuSans-115\"/>\n     <use x=\"216.210938\" xlink:href=\"#DejaVuSans-111\"/>\n     <use x=\"277.392578\" xlink:href=\"#DejaVuSans-114\"/>\n     <use x=\"318.505859\" xlink:href=\"#DejaVuSans-40\"/>\n     <use x=\"357.519531\" xlink:href=\"#DejaVuSans-49\"/>\n     <use x=\"421.142578\" xlink:href=\"#DejaVuSans-41\"/>\n    </g>\n   </g>\n  </g>\n </g>\n <defs>\n  <clipPath id=\"pa3021f6c68\">\n   <rect height=\"81.6\" width=\"81.6\" x=\"7.2\" y=\"19.27875\"/>\n  </clipPath>\n  <clipPath id=\"pd96658de3c\">\n   <rect height=\"81.6\" width=\"81.6\" x=\"99.6\" y=\"19.27875\"/>\n  </clipPath>\n  <clipPath id=\"pd8fe64a99a\">\n   <rect height=\"81.6\" width=\"81.6\" x=\"192\" y=\"19.27875\"/>\n  </clipPath>\n </defs>\n</svg>\n",
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAARgAAABsCAYAAACxSq3vAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAKUklEQVR4nO3dW2hU2xkH8P83GTU1MyaNCTpHDUkbL6PkFFTUg1GPtiYi7Snig2iV1AteSlulPrSC4K1YqIKgD4cD2ocTS0XtS6VqUrwU9YjSKjFSc4SiRiuKtUmMxViTrD7MHDtaJ85lr7X25f+DwcnMntmf397zzfrW7NkjSikQEekQsh0AEfkXCwwRacMCQ0TasMAQkTYsMESkDQsMEWljrcCISKWIzNO8DhGRxuT1VSJyR0QOp9y/UUSm6YzBLSzk+wMROSsiX4jId5K3BSbfAPdxwO4IphKAtuSLSAjAtwFcTd70BwDz31rscwA/0hWDy1TCbL5/AWArgLrkv0Cw8g1wH0fY1ooBrAUwU0Q+AvBnJDZEP4BVyft/A+BfAKoAfB9AIRLJegngT0qp3SLycwCfJG/7oVKqXURaALQCuAkglnwMlFL/FJFIagBKqY7kO60o/x9xaDTfAD4EsFEppUSkW0SiAcs3wH0cUEpZuQD4GMAvAdQA+Cx5WxzAZ0hU/utIjLB+AGAjgDVIJBgABMBIAE3Jv2sBfJq8/hRAUfL6SQDDU9ZZCeDwW3EcAfCBrTz4Nd8ALqSs+zCAiiDl20bOk3+7ah93wyRvHMDHInIewKcAhiVv/5tSqh/APwCUADgK4EMR+S2ABUgk8kZy2b8AqE5e/1Ip9e/kdclg/QIgCO+mXzGV776U68MAdKYsE6R8AwHex222SK8AFAD4EkCzUuonACAigwCMwpsJEQCvlFI/E5HBAC4B+B6AbyXvnwrg78nr/SmPu43ERno6QBylAB7n9T/xBtP5vpFsDW4AGKaUepZcJij5BriPWy0wNwH8CsA3kNgZzyOR8N8BaH7H8p+IyI8BDEViCPhIRM6JyBcA/gOg4R2P+SOAjwD8VUS+i8TE4zdF5PdKqcUi8nUAD5PvIn5nNN8Afo3E3MDXAGwDgIDlG+A+Dkn2aL4kIgLgc6XUijT3/xTAFaXUFbOR+RPzbZ7bc+7rAkNEdrlhkpeIfIoFhoi0YYEhIm0G/BRJRDhB4wClVCbHKgBgzp2Sac6Zb2ekyzdHMESkDQsMEWnDAkNE2rDAEJE2LDBEpA0LDBFpwwJDRNqwwBCRNiwwRKQNCwyRR1y5ciWj01TW1dUhHA4jFLL/8h7wdA08jNoZ/KqAeX77qkBhYSFevHiR8fIrV65EW1sbWlpasnpcrvhVASIPKyoqymr5LVu2oL6+HtFoVFNEmfF0gYlGo5g4cSJmzJiR8VnOb9++bTtsXxs6dGjWZ55vampCeXm57dBdLdt2Z9y4cdi+fTsqKio0RZSh9/zsgnLz5fz58yoXpuPM5sVmO6f5Xnbt2pXTNnF6u/gt37FYzBV5zTbfnp6DGSj2gSROY2qOCtAcTK7bBHB2u2Sac6/ku6KiAvfu3cvpsSb293T59nSLRObV1taiubkZ3d3d6UYEOUv3Ltja2orly5ejoKDAof+Fd2zatAlKqZyLi20cwRjgpxFMvkUkH/F4HG1tbRkt65cRjBP55giGKAPxeNx2CJQlzxWYxYsXOzIcp/9XXFyMuXPnYvXq1bhw4YLjLVC+Ro0aZXX9lD3PtUheGTKm8kqLtHXrVuzatcvW6t+rq6sLJSUlGS3LFul/bLZILDAGeKXAPHz4ELFYzNbqM5LptmOBSZgyZQquXbvmUDTpcQ6GXguFQohEIigrK8PevXtftz9uLy4AUFpaiiFDhtgOw/VEBCJipLgMhAUmgCKRCObPn4/Nmzdj8+bNtsPJyqpVqzBhwgTjo1DKDQtMAEUiEcyePRvr1q2zHUrWli1bhrFjx9oOgzLEORgD3DAHs2/fPmzatEnHU1sxcuRIPH78OO39QZ+Dccs+zhFMQPipuABAQ0OD7RAoAyww5EmLFi2yHQJlgAXGxw4dOuSKA+R0qK6uth2CNm45sNEJnIMxwNYcjB920IEMtB29PAfjp32cIxgi0oYFxmfWr1/vm+E1eR9bJANMtkhBKixskdJzyz7OEQwRacMCQ0TasMAQkTYsMGTEpEmTEIvFMGjQIITDYVRXV6O+vt52WKQZJ3kN4CTvmzk/cuQIlixZ4uhzvs1rk7xObze37OMcwZBxThSX48ePOxCJP7npe2dh2wGYEgqFXPvu7ieFhYXo6+tDX18fAH0jqqtXr2p5Xi9z4zlyAlNgWFzMePnypZH13L9/38h6KD9skciTLl68aDsEykBgRjCkTyQSQW9vL169eqV9XW5sAyg9jmAob8+fP0dPT8/reZd3iUajOHv2LFvVgGGBISNisRjmzp1rOwwyjC0SOWrmzJl48uQJ2tvb0dPT8/r2aDRqMSp3Gjx4MMrKynydm8AcaGezdw/qgXapOXcqLj/98JrXD65LxQPtyPPcVDzdZuHChbZDeCdPtEiVlZVYunQppk2bZjsUVyopKcGZM2cwefJk26G8wemC4ObfzbbF7Z+qeaLA3Llzx3YIrtbR0WE7BCNaW1tth0BZ8sQcjBe/4JhK9xxMUFqHsrIyPH36NKNlgzIH45YRTLp8e2IEQ/bEYjE8evTonfcVFxdjzJgxiEQiWLNmDVavXq0lBre8iJxQVFSEeDyOyspK26EYwQJDA3rx4kXa+zo7Ow1G4g9tbW0YPXq07TCMYYtkgJdbpHA4/M4jdEUE/f392tb7lf7+fhQUFGT9OLe2SH76aDoVWyTKSW9vr9X1Hzt2zOr6KT88DoZcrampyXYIlAe2SAZ4uUWyrbS0NKeP4dkimcUWiTzFLS8cJwwfPhzLli1DXV2d7VCMY4Eh12lvb7cdgqNOnz6NqVOnOv68R48edfw5ncYWyQC2SNmpqqrC3bt383oON7VITm+fPXv24ObNm2hubk57jJJp6fLNAmMAC0xmnNxGfi4wbmwf+W1qIjIuEAVmzpw5tkOg9xg/frztEEgDX0/yunEoSdwuQRKIEQwR2cEC4wMnTpywHULGnj17ZjsET3PrmevSYYHxAS8dTr9t2zbbIXiSiEBEcOrUKduhZIUFxgcuX75sO4SM8RcZg4UFxgdaWlpQU1ODhoYG26G8obOzE42NjZg1axaqqqpQVlaG69ev2w7LuOnTp+f1eE+fd0cplfYCQLnhkivbcafEP2Cencz53bt3c86XE2znOtucm4hlxIgRaufOnaqtrS2nnG7YsMF6PnPNN0cwPrNjxw7bIdBburu7cfLkSezevTurx128eBH79+/HpUuXNEWmn6+/KuCW4y2UwR9eC4fDqK6uRjweR3l5OVasWIHa2tp8njKtjo4OHDhwAK2trXjw4AFu3bqFrq4uLevKVqY5N7WPh0IhhEIh1NbW4ty5c+9dvrGxEWvXrkV/fz96e3uNnD0wH+nyzQJjgMkCQwluKzB+ly7fbJGISBtPFJgFCxZktNzBgwcxb9481NTU+PoHxYm8whMtktexRTKPLZJZbJGIyDgWGCLShgWGiLRhgSEibVhgiEgbFhgi0mbAj6mJiPLBEQwRacMCQ0TasMAQkTYsMESkDQsMEWnDAkNE2vwXMr9A2HMaSacAAAAASUVORK5CYII=\n"
     },
     "metadata": {
      "needs_background": "light"
     }
    }
   ],
   "source": [
    "#################################################################################\n",
    "#########                  check the evaluation results                  ########\n",
    "#################################################################################\n",
    "print(\"Predictions:\", predicted[:3].detach().numpy())\n",
    "print(\"Ground Truth:\", test_labels[:3])\n",
    "evaluate_fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(4, 4))\n",
    "axes = axes.flatten()\n",
    "\n",
    "for i in range(3):\n",
    "    a = axes[i]\n",
    "    a.imshow(test_imgs[i,0,:,:], cmap='gray', vmin=0, vmax=1)\n",
    "    a.set_title(test_labels[i], fontsize=8)\n",
    "    a.axis('off')\n",
    "    \n",
    "evaluate_fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "#################################################################################\n",
    "#########                     save the model as ONNX                     ########\n",
    "#################################################################################\n",
    "# store as ONNX\n",
    "# example input\n",
    "x = next(iter(train_loader))[0]\n",
    "torch.onnx.export(model, x, os.path.join(output_path,'mnist_model.onnx'), opset_version=11,\n",
    "                  export_params=True, input_names=['input'], output_names=['output'],\n",
    "                  dynamic_axes={'input': {0: 'batch_size'},\n",
    "                                'output': {0: 'batch_size'}})"
   ]
  }
 ]
}