mnist_cnn_snn.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import numpy as np\n",
    "import cv2\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing MNIST dataset and creating train and test loder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_epochs = 5\n",
    "batch_size_train = 64\n",
    "batch_size_test = 1000\n",
    "learning_rate = 0.01\n",
    "momentum = 0.5\n",
    "log_interval = 10\n",
    "\n",
    "random_seed = 1\n",
    "torch.backends.cudnn.enabled = True\n",
    "torch.manual_seed(random_seed)\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "  #torchvision.datasets.MNIST('/files/', train=True, download=True,\n",
    "  torchvision.datasets.MNIST('.', train=True, download=True,\n",
    "\n",
    "                             transform=torchvision.transforms.Compose([\n",
    "                               torchvision.transforms.ToTensor()\n",
    "                             ])),\n",
    "  batch_size=batch_size_train,pin_memory = True, shuffle=True)\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "  #torchvision.datasets.MNIST('/files/', train=False, download=True,\n",
    "  torchvision.datasets.MNIST('.', train=False, download=True,\n",
    "\n",
    "                             transform=torchvision.transforms.Compose([\n",
    "                               torchvision.transforms.ToTensor()\n",
    "                             ])),\n",
    "  batch_size=batch_size_test, pin_memory = True,shuffle=True)\n",
    "\n",
    "#torchvision.transforms.Normalize(                                 (0.1307,), (0.3081,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAZQAAAELCAYAAAD+9XA2AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4xLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy8li6FKAAAem0lEQVR4nO3deZRU1bn38d8DKjJFQBAQGUSccAkmTuCA14gKxOkqEQ0Lg1GXBscI8caBOPsajeEaTdCl6+J4yVUJosbpYkQFBxJYgkLUAIGAYWzmQQbZ7x9VnPfs83YVVdW7uk53fz9r9Vr76X2GXfSmnjp7n9rHnHMCAKCmGlW6AQCA+oGEAgAIgoQCAAiChAIACIKEAgAIgoQCAAiiXicUM1toZv0reP4lZvZvlTo/SkffQakact+pUUIxs4vM7BMz22RmK7LlEWZmoRpYDmb2hpltzP5sN7NtsfixEo/5nJndEbCNo2Nt2mhmW8zsWzNrHeoclUTf8Y4ZtO9kj7mfmY03s3VmtsbMngl5/Eqi73jHTNX7TskJxcxGSnpY0oOSOkhqL+kqSSdK2ivHPo1LPV9IzrmBzrkWzrkWkp6X9MCu2Dl3VXJ7M9ujAm28O9amFpIekvSOc25NbbclNPpOrZgkabGkzpL2kzSmQu0Iir5T9jbW7H3HOVf0j6R9JG2SdMFutntK0lhJr2e375/d9xlJKyUtknSbpEbZ7e+Q9Fxs/26SnKQ9svEUSXdLmiZpg6S3JbWNbT8se8wqSbdKWiipfwFtvCfxu/7ZfW+RtEzSOEmXS5oS22aPbNu6SRohabukbZI2SpqY3WaJpBslfSZpnaTxkpqU8O9t2dc1tJS/V5p+6Dvl7zuSBkmav+vfpr780HfS/75T6hVKX0lNlPkUtDs/knSvpJaSpkp6RJk/bndJp0i6RNKlRZz7R9nt91PmE8koSTKznsp0omGS9pe0r6QDijhu0gGSWkjqoswfLifn3O8l/Y+k+1wms/97rPpCSacr83qPzrZPZtbYzNaaWZ8C2nKqpNaSJhb9KtKHvhNTpr7TR9KXkp4zsyozm25mJ9Xg9aQFfScmje87pSaUtpJWOed27PqFmX2YbegWM+sX23aSc26ac26nMtl0iKSbnXMbnHMLlbmkGlbEucc5575yzm2R9IKko7K/HyzpNefc+865rZJGS9pZ4uuTpB2S7nDObcueq1T/6Zxb5pyrkvTarvY65751zrVyzn1cwDF+LOkF59zmGrQjLeg7hSu17xwgaaAyn6Q7KDNE9IqZtalBW9KAvlO4irzvlJpQqiS1jY/xOedOcM61ytbFj7s4Vm6rTHZfFPvdIkmdijj3slh5szLZXMp8OojO5ZzblG1LqZY757bVYP9dcrW3IGbWXNIFkp4O0JY0oO8UrtS+s0XSPOfcU8657c655yUtV+YTfl1G3ylcRd53Sk0oH0naKuncAraNL2e8SplPC11jv+si6etseZOkZrG6DkW0aakyE5CSJDNrpszlZ6mSyzDvrm3lWrZ5sDJvBlPLdPzaRt8pf9+ZXYZjpgF9J+XvOyUlFOfcWkl3Svq9mQ02sxZm1sjMjpLUPM9+3ypzuXivmbU0s67KTB49l93kU0n9zKyLme0j6eYimvWSpLPM7CQz20vSXQr7PZtZknqZ2ZFm1lTS7Yn65cqMV4b2Y0lPu+wsWV1H36mVvjNBUnszG5odMx8iqZ0yb8h1Fn0n/e87Jb9w59wDyvxRbpK0QpkX9rik/5D0YZ5dr1Um6y5QJvv9t6T/yh7zf5WZZJotaYYyY3+FtmeOpKuzx1sqaY0ydzsE4ZybK+k+Ze74+FLS+4lNnpTUO3vP/0u7O172P/pGM8s5DGFmXST1k/RsyQ1PIfpOefuOc26VMp/ib1bmLp9Rks5xzq0u/VWkA30n3e87Vk8++AIAKqxeL70CAKg9JBQAQBAkFABAECQUAEAQJBQAQBBFrWZpZtwSlkLOubQv202/SadVzrl2lW5EPvSd1Kq273CFAjRci3a/CVCtavsOCQUAEAQJBQAQBAkFABAECQUAEAQJBQAQBAkFABAECQUAEAQJBQAQBAkFABAECQUAEAQJBQAQBAkFABBEUasNA/XRFVdc4cVz5syJygsWLPDqli1bVittAuoirlAAAEGQUAAAQTDkhXrpqquu8uK+fftG5ZYtW3p1Z599thfv2LGj2rIk9e/f34s/+eSTGrUTqE+4QgEABEFCAQAEQUIBAATBHArqpe7du3vx6aefHpXbt2+fd99PP/00Ki9evNiru/zyy72YOZS67+233/bi448/3ot79OgRlVeuXFkrbaqruEIBAARBQgEABFG2Ia+xY8d6cXxo4KmnnirXadFAHXHEEV58ySWXePG+++4blV944QWv7t577/XiRYsWReVvvvnGq9t7771r1E6kT7du3bw4eVv55MmTo3Lv3r1ro0l1FlcoAIAgSCgAgCBIKACAIMw5V/jGZgVvnDzuihUronL8Fk5Jmj17dsFtSLv4WP6wYcO8ul/96ldevGbNmiDndM5ZkAOVSTH9phjx5VXuuecer65169ZeHJ83Sf5dksurNCAznHPHVLoR+ZSr78T9+te/9uIbb7wx57bz58/34scee8yLJ02aFKRN//znP71427ZtQY4bULV9hysUAEAQJBQAQBAkFABAEGX7Hsq6deu8uG3btlF5yJAhXt28efOi8ubNm8vVpGDatGkTlS+++GKv7vbbb4/K8e8+SFKHDh28ePjw4eEb14DEl6RPzpkkxb9r0oDnTFCNqqqqvPVbt26Nyp07d/bqHnzwwbxxqUaOHOnFY8aMCXLccuMKBQAQBAkFABBE2W4b/sEPfuDFr7zySs5tJ0yYEJXvv/9+r27ZsmVe/K9//avQJhSlS5cuUfm4447z6gYOHOjFp5xySlQ+8MADCz7HggULvPjggw8upok5NdTbhuPLouy5555eXXJ5lfhSLNu3by9Hc+oibhuWdMYZZ3jxm2++6cWXXXZZVJ45c6ZXd84553jxF198EZU3bNiQ97xm/++/7R/+8AevLvm+d+ihh+Y9VgVw2zAAoHxIKACAIEgoAIAgynbb8FtvveXF8XHJM88806u74IILonJy7iU53h1fgiA5n9KkSRMvfvHFF3O2L3mbaXyMvUWLFjn3q4mXX365LMdtKOJPzpP8Meik5JL0aZg3adeuXVR+8sknvbq5c+dG5S1btnh1zzzzjBcvXLgwfOMasBNOOMGLV69e7cXjxo3Lue+sWbOCtCF5K/v48eODHLe2cYUCAAiChAIACIKEAgAIomxzKMkxwfjcyN133+3VXXnllVE5uVxJvkeuJrdNuvXWW3fbzupMnDjRi08++WQvji8jk/Ttt99G5Ztvvtmre+KJJ0pqDzKS/5577JG7+y5durTczSnaddddF5X79Onj1Z111lk597vooou8eMCAAVE5ucw5aq6Y7+aV65yVaEMIXKEAAIIgoQAAgijbkFc+o0eP9uI//elPUTl5eR+/nVfyV/5ctWqVV9ezZ08vjg8/JSVvBXznnXei8qWXXurV5buNODm0d+qpp0blDz/8MOd+KF58eRxJ+vrrr6NyciXnNDjssMO8OH57anLYNH47fPIJgrfccosXx1frDrW6bUM2ZcoUL04uvVQu8SWc9tlnn1o5Z7lxhQIACIKEAgAIgoQCAAiiInMoSR9//HG1ZUm64YYbcu6XXDIhOcaeXEIhbvLkyV78s5/9LConl7PO55e//KUXM29SPv379/fi+fPnV6gl1TvkkEO8ODk2H196Jek3v/lNVH7ooYe8ussvv9yL4/OKzz//vFdXrsc71GfJv1MyLpdmzZpF5caNG9fKOcuNKxQAQBAkFABAECQUAEAQqZhDKVVyvqKY+Yvk0ufJR3nmU1VVFZXHjh1b8H6omfg8gyRdf/31UTn592zfvr0Xx/9mITVv3jwqd+rUyavLN2cye/ZsL3788cej8tq1a7265ONh40u4JL8zlVy2H+kV/85afcEVCgAgCBIKACCIOj3kVROHH364F5900kk5t00OQZx//vlRef369WEbhpxGjhzpxaeffnpUTv49kytNjxo1KiqHXIm4adOmUXl3w6bxJy/efvvtXl181eD4MJok9e3bN+cx8618jfJr1aqVF8f/HvPmzcu7b8eOHaNycsg2eTt4XcEVCgAgCBIKACAIEgoAIIgGO4eS72mOGzdu9OLkePfUqVPL0iYUp1evXlF50aJFXl3yMQjHHntsVL7wwgu9uuRyJStWrCi4DfFHKCTHveO390r+PF2+ZYE2bdrkxR999JEXx1/LsGHDvLr4EkIovyeffNKLBw4cGJX/+Mc/enXJOD7fknxC43nnnefF8Xm1o48+2qs78cQTo3LykR7JpaveeOMNlRNXKACAIEgoAIAgSCgAgCAazBxK69atvfiHP/xhzm3vv/9+L3700UfL0iaEM2jQIC9+8803vfiggw6KyjNmzPDqFi5c6MXvvvtuSW3Y3XdCunfvHpVfffVVry7fUvzJR1vHjRkzpsDWoRxuu+02L95zzz2j8tChQ726ZJzPAw88kLNuw4YNXhyfm9l33329uh49ehR8zhC4QgEABEFCAQAE0WCGvG666SYvzveEtJ07d5a7OQhszpw5XjxgwAAvjt/C+5Of/MSr69atmxcnV/Ath379+uWN46ZPn+7Fy5cvj8rJoT3Uri+++MKLhwwZEpWT7zmDBw/24vgTPpO3kSefBBtfnTq5cnaanlzKFQoAIAgSCgAgCBIKACAIS37lP+/GZoVvnALf+973ovInn3zi1TVqlDuXXnPNNV6c9qcyOuds91tVTtr6TXKOpE+fPl4cHwcvRnx5ekk67rjjvPiwww6Lyi+99JJXF3+6Y3Kplfh+kjR8+PCS2leNGc65Y0IdrBzS1ndCevbZZ6Ny8vbefI8sSIlq+w5XKACAIEgoAIAg6vVtw/Hb6RYsWODV5fsG6axZs8rWJlTeuHHj8sZXXnllbTYH+P+e2FhXcYUCAAiChAIACIKEAgAIol7PoWzevLnacnW2bt0alT///POytQkAkor5+kaacYUCAAiChAIACIKEAgAIol7PofTu3Tsq9+rVK++2EydOjMrr168vW5sAICn+XiVJZ599thcnn/CZVlyhAACCIKEAAIKo10NexRg/fnylmwCggdp77729OL5SusSQFwCggSGhAACCIKEAAIKo13Mo//jHP6Jy8omNRx55pBcvWbKkVtoEAPUVVygAgCBIKACAIEgoAIAgrJhlk82szq6x3KZNGy9u166dF3/55Ze12ZygnHOpfn5oXe439dwM59wxlW5EPvSd1Kq273CFAgAIgoQCAAiiXt82HLd69eq8MQCgZrhCAQAEQUIBAARBQgEABFHsHMoqSYvK0RCUrGulG1AA+k060XdQqmr7TlHfQwEAIBeGvAAAQZBQAABBkFAAAEGQUAAAQZBQAABBkFAAAEGQUAAAQZBQAABBkFAAAEGQUAAAQZBQAABBkFAAAEGQUAAAQdTrhGJmC82sfwXPv8TM/q1S50fp6DsoVUPuOzVKKGZ2kZl9YmabzGxFtjzCzCxUA8vBzN4ws43Zn+1mti0WP1biMZ8zszsCtrG/me2MtWujmQ0NdfxKo+94xwzad7LH3M/MxpvZOjNbY2bPhDx+JdF3vGOGft8ZnXjP2WJm35pZ60L2LzmhmNlISQ9LelBSB0ntJV0l6URJe+XYp3Gp5wvJOTfQOdfCOddC0vOSHtgVO+euSm5vZsU+iCyUf8ba1cI593yF2hEUfadWTJK0WFJnSftJGlOhdgRF3yl7G++Ov+dIekjSO865NYUeoOgfSftI2iTpgt1s95SksZJez27fP7vvM5JWKvMkttskNcpuf4ek52L7d5PkJO2RjadIulvSNEkbJL0tqW1s+2HZY1ZJulXSQkn9C2jjPYnf9c/ue4ukZZLGSbpc0pTYNntk29ZN0ghJ2yVtk7RR0sTsNksk3SjpM0nrJI2X1KTAf+P+khaW8vdJ8w99p1b6ziBJ83f929SXH/pO+ftOoj2WfV1DC92n1CuUvpKaKPMpaHd+JOleSS0lTZX0iDJ/3O6STpF0iaRLizj3j7Lb76fMJ5JRkmRmPZXpRMMk7S9pX0kHFHHcpAMktZDURZk/XE7Oud9L+h9J97lMZv/3WPWFkk5X5vUenW2fzKyxma01sz55Dt3RzJab2QIze8jMmtXg9aQFfSemTH2nj6QvJT1nZlVmNt3MTqrB60kL+k5MGd93djlVUmtJEwttfKkJpa2kVc65Hbt+YWYfZhu6xcz6xbad5Jyb5pzbqUw2HSLpZufcBufcQmUuqYYVce5xzrmvnHNbJL0g6ajs7wdLes05975zbquk0ZJ2lvj6JGmHpDucc9uy5yrVfzrnljnnqiS9tqu9zrlvnXOtnHMf59hvTnbbjsp0jD7KXObXdfSdwpXadw6QNFCZT9IdlBkiesXM2tSgLWlA3ylcqX0n7seSXnDObS70pKUmlCpJbeNjfM65E5xzrbJ18eMujpXbKpPdF8V+t0hSpyLOvSxW3qxMNpcynw6icznnNmXbUqrlzrltNdh/l1ztzcs5t9Q59zfn3E7n3HxJ/6FM563r6DuFK6nvSNoiaZ5z7inn3HaXmXtbrswn/LqMvlO4UvuOJMnMmku6QNLTxexXakL5SNJWSecWsK2LlVcp82mha+x3XSR9nS1vkhQf1ulQRJuWKjMBKUnKDg/tW8T+SS4R765tye1Dc8qMadZ19J3y953ZZThmGtB3au99Z7AyH0KmFrNTSQnFObdW0p2Sfm9mg82shZk1MrOjJDXPs9+3ylwu3mtmLc2sqzKTR89lN/lUUj8z62Jm+0i6uYhmvSTpLDM7ycz2knSXwn7PZpakXmZ2pJk1lXR7on65MuOVQZjZqWbWOVvuIun/qLCx41Sj75S/70iaIKm9mQ3NjpkPkdROmTfkOou+Uyt9Z5cfS3raZWfnC1XyC3fOPaDMH+UmSSuUeWGPKzM082GeXa9VJusuUCb7/bek/8oe83+VmWSaLWmGMmN/hbZnjqSrs8dbKmmNMnc7BOGcmyvpPmXu+PhS0vuJTZ6U1Dt7z/9Luzte9j/6RjPLNQxxjKSPzWyzMv9OMyX9rNT2pwl9p7x9xzm3SplP8Tcrc5fPKEnnOOdWl/4q0oG+U/b3nV0fYPtJerbY9lqRCQgAgGrV66VXAAC1h4QCAAiChAIACIKEAgAIgoQCAAiiqNUszYxbwlLIOZfqLzzSb1JrlXOuXaUbkQ99J7Wq7TtcoQAN16LdbwJUq9q+Q0IBAARBQgEABEFCAQAEQUIBAARBQgEABEFCAQAEQUIBAARBQgEABFHUN+UB+H76059G5UcffdSre+SRR7z4hhtuqJU2AZXCFQoAIAgSCgAgCBIKACCIop4pz8qf6cRqw7WnR48eXvzuu+9G5Y4dO3p127dv9+IBAwZE5ffee68MrSvaDOfcMZVuRD71qe/UM9X2Ha5QAABBkFAAAEFw23AJDjjggKh82mmneXVHHXVUzv0GDx7sxZ06dYrKmzZt8uqOP/54L547d27R7UR45557rhfvv//+UTk5fLznnnt6cbt2qX6WFVBjXKEAAIIgoQAAgiChAACCqFdzKF27dvXiQYMGReX4WLck9erVy4u/+93vRmUz/y7c5Nh48+bNo3Lr1q1La2xC/JiS1L59ey9mDqUyWrVq5cUjRoyoUEuA9OMKBQAQBAkFABBEnR7yGjlypBcPGzbMi5PDWuWwdetWL/7qq6+icvfu3b26P//5z1689957R+XZs2d7dTNnzgzVRNTA/fff78XJYdV83n//fS9+++23g7QJYTVr1syL48PfSfH/s5J0zDH+l8UPPvjgqHzIIYfkPW/8vSJp3bp1XnznnXdG5fXr1+c9biVxhQIACIKEAgAIgoQCAAiizq02HL+dNjkG2bJly4KPs3jxYi/u3LlzVJ41a5ZX98orr3jx559/HpU/+ugjr27JkiUFtyEUVhsOK96PZsyY4dUddNBBXhy/xTz5f6lDhw5evHLlylBNDKXBrDbcr18/L7711lujcvJveuCBBybbEJWLeb9M2rFjhxevXbs2Ku+1115e3Xe+8x0vnjx5clQ+88wzS25DQKw2DAAoHxIKACAIEgoAIIg69z2UY489Nirvbs7kiSeeiMrjxo3z6uLzIJI/hplcSj75XRPUb/Hx9eR3ifKNoU+ZMsWL42PkqH3x+dbx48d7dcn5rXwmTZoUlSdMmODVFfOdkNWrV3vx1KlTo3LysRfTpk3z4v79+xd8nkriCgUAEAQJBQAQRJ0b8mrSpEnOup07d3px/PL0448/LlubUL8MHz684G3jtwJfffXVXt327dtDNQklWL58eVS++OKLvbr4UFXyKwRJVVVVYRtWjeQyUcklXv7yl7+UvQ0hcIUCAAiChAIACIKEAgAIIvVzKMk5k9/+9rc5t12zZo0Xs1w4CtGnTx8vTi57kU98mZ4vvvgiWJsQVvJRAmkQnye5/vrr826bfIxCWnGFAgAIgoQCAAiChAIACCL1cyjJJRI6duyYc9trr7223M1BPdCmTRsvHjNmjBcnlxLPJ76cxkMPPeTVnXHGGV781ltvReX77rsv53HQMAwYMCAqJx87vHTpUi9OPiYjrbhCAQAEQUIBAASR+iGvgQMH5qxLDhPMmzfPi5s2bRqVt2zZErZhqLN69uzpxccdd1zJx7rooouicnLpn3znTQ7dDh06tOQ2oG76xS9+EZWTq1h/8MEHXhxfRibNuEIBAARBQgEABEFCAQAEkfo5lHySt39Onz7di2fPnh2VR48e7dW9+uqr5WsYUi3+REYp/1MYdyc+b1LMcYYMGeLFL774YlR++eWXS24P6o58S/xMnDixFlsSDlcoAIAgSCgAgCBIKACAIFI/h5K8Hzs+L5J8bGZSvH7SpEle3WeffebF8UeCTps2zau74447vPibb77Je16kW3JJlJrMoYSSb0kh1A/dunXz4nbt2uXcdvLkyWVuTXlwhQIACIKEAgAIIvVDXnPmzPHivn37RuXk6q7JFTuPOOKIqNyiRQuv7sgjj8x5zhNPPNGLu3Tp4sWXXXZZVGZJl7qhU6dOlW4CGrhkH0x+7aFQ7du39+LOnTtH5b/+9a9e3VlnneXFr732WknnLBRXKACAIEgoAIAgSCgAgCBSP4eSFJ+zGDFiRN5tDzvssKjcqlUrr+68887z4vhSGF27dvXq4kuUS1KjRo1y1iGdzjnnnLIcd/78+VH5/fff9+qGDx9elnMivYYNGxaVDz/8cK+uSZMmXmxmOY+zcuXKnHXJ/eK3vf/tb3/z6uJL+kjMoQAA6ggSCgAgCBIKACAIK2bZCTOr/BoVZRJfFuHuu+/26pKPZ920aVNUbtmyZVnbVQjnXO7B2BRIQ7+5+uqro/Ijjzzi1dVk6ZX4fNruHgEct2bNGi9u27ZtyW2ogRnOuWMqceJCpaHv5PO73/3Oi6+44oqo3LhxY68u39zHtm3bvLrFixd78YQJE6LyihUrvLrXX389Kn/99dde3caNG3O2vYaq7TtcoQAAgiChAACCqHO3DZfLwoULo3LysjEpubwB0m/evHlROTnEVYknNnK7ef3w9NNPe/GCBQui8t///nev7pprrvHi0047LSqPGjXKq0sOpdUVXKEAAIIgoQAAgiChAACCSN0cSnLp+KOPPtqLH3/88ai8devWks/TrFkzLx45cmRUvummm/Lu++mnn5Z8XlTGW2+9Vekm6OGHH47KU6ZMqVxDEMz06dPzxnHJeZKqqqqo/MQTT4RtWIVwhQIACIKEAgAIIhVDXr17947K48aN8+p69Ojhxd///vejcnJoKv5td8lfYfj444/36gYNGuTFhx56aFROfqN1yZIlXnzXXXcJdVeyj4VaFTi5QuzPf/5zL44Pc+3YsSPIOVF35fumfF3FFQoAIAgSCgAgCBIKACCIVMyhxFdaTc6ZJMWfvDdw4ECvLrm6Z3wl2GKsW7fOi5966ikvTq4Ui7olvvKw5D91UZJuueWWqNy0adO8x7rnnnuicvLWz+TcGxq2k08+2YvzPZWxruIKBQAQBAkFABAECQUAEEQqnth40EEHReUPPvjAq+vQoUM5Tqn169d78cyZM6Pys88+69Ulv7eQNjyxESXiiY21KPlEz/gcSvv27Wu7OTXFExsBAOVDQgEABJGK24bjt21ed911Xt25557rxf369YvKixYtynmcZP17772Xs07yn7QGAKFNnjzZi+NLTtUXXKEAAIIgoQAAgiChAACCSMUcStxLL72UNwaAuuizzz7z4r59+0blnj17enVz586tlTaFxhUKACAIEgoAIAgSCgAgiNTNoQBAfTR27FgvPv/886PysmXLars5ZcEVCgAgCBIKACAIhrwAoBbMmzfPiw888MAKtaR8uEIBAARBQgEABEFCAQAEUewcyipJi3a7FWpT10o3oAD0m3Si76BU1fadoh4BDABALgx5AQCCIKEAAIIgoQAAgiChAACCIKEAAIIgoQAAgiChAACCIKEAAIIgoQAAgvi/h3DkgbZEx4IAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "examples = enumerate(test_loader)\n",
    "batch_idx, (example_data, example_targets) = next(examples)\n",
    "example_data.shape\n",
    "\n",
    "fig = plt.figure()\n",
    "for i in range(6):\n",
    "    plt.subplot(2,3,i+1)\n",
    "    plt.tight_layout()\n",
    "    plt.imshow(example_data[i][0], cmap='gray', interpolation='none')\n",
    "    plt.title(\"Ground Truth: {}\".format(example_targets[i]))\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.mnist = nn.Sequential(\n",
    "            nn.Conv2d(1, 8, kernel_size=3),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(8, 16, kernel_size=3),\n",
    "            nn.BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),\n",
    "            nn.MaxPool2d(2, stride=2),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout2d(),\n",
    "            nn.Conv2d(16, 20, kernel_size=5),\n",
    "            nn.BatchNorm2d(20, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),\n",
    "            nn.MaxPool2d(2, stride=2),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout2d(),\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(320, 50),\n",
    "            nn.BatchNorm1d(50, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(50, 10),\n",
    "            nn.LogSoftmax()\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.mnist(x)\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "network_cpu = Net()\n",
    "network = network_cpu.cuda()\n",
    "optimizer = optim.SGD(network.parameters(), lr=learning_rate,\n",
    "                      momentum=momentum)\n",
    "train_losses = []\n",
    "train_counter = []\n",
    "test_losses = []\n",
    "test_counter = [i*len(train_loader.dataset) for i in range(n_epochs + 1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epoch):\n",
    "    network.train()\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        optimizer.zero_grad()\n",
    "        output = network(data.cuda())\n",
    "        loss = F.nll_loss(output, target.cuda())\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        if batch_idx % log_interval == 0:\n",
    "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "                epoch, batch_idx * len(data), len(train_loader.dataset),\n",
    "                100. * batch_idx / len(train_loader), loss.item()))\n",
    "            train_losses.append(loss.item())\n",
    "            train_counter.append(\n",
    "                (batch_idx*64) + ((epoch-1)*len(train_loader.dataset)))\n",
    "            torch.save(network.state_dict(), 'model.pth')\n",
    "            torch.save(optimizer.state_dict(), 'optimizer.pth')\n",
    "def test():\n",
    "    network.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            output = network(data.cuda())\n",
    "            test_loss += F.nll_loss(output, target.cuda(), size_average=False).item()\n",
    "            #test_loss += F.nll_loss(output, target, size_average=False).item()\n",
    "\n",
    "            pred = output.data.max(1, keepdim=True)[1]\n",
    "            #correct += pred.eq(target.data.view_as(pred)).sum()\n",
    "            correct += pred.cpu().eq(target.data.view_as(pred.cpu())).sum()\n",
    "            test_loss /= len(test_loader.dataset)\n",
    "            test_losses.append(test_loss)\n",
    "            print('\\nTest set: Avg. 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": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\crsum\\Anaconda3\\envs\\pytorch\\lib\\site-packages\\torch\\nn\\modules\\container.py:92: UserWarning: Implicit dimension choice for log_softmax has been deprecated. Change the call to include dim=X as an argument.\n",
      "  input = module(input)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 1 [0/60000 (0%)]\tLoss: 2.383272\n",
      "Train Epoch: 1 [640/60000 (1%)]\tLoss: 2.304424\n",
      "Train Epoch: 1 [1280/60000 (2%)]\tLoss: 2.208299\n",
      "Train Epoch: 1 [1920/60000 (3%)]\tLoss: 2.146656\n",
      "Train Epoch: 1 [2560/60000 (4%)]\tLoss: 2.014208\n",
      "Train Epoch: 1 [3200/60000 (5%)]\tLoss: 1.935144\n",
      "Train Epoch: 1 [3840/60000 (6%)]\tLoss: 1.930165\n",
      "Train Epoch: 1 [4480/60000 (7%)]\tLoss: 1.761986\n",
      "Train Epoch: 1 [5120/60000 (9%)]\tLoss: 1.739454\n",
      "Train Epoch: 1 [5760/60000 (10%)]\tLoss: 1.647626\n",
      "Train Epoch: 1 [6400/60000 (11%)]\tLoss: 1.554630\n",
      "Train Epoch: 1 [7040/60000 (12%)]\tLoss: 1.441419\n",
      "Train Epoch: 1 [7680/60000 (13%)]\tLoss: 1.401727\n",
      "Train Epoch: 1 [8320/60000 (14%)]\tLoss: 1.434840\n",
      "Train Epoch: 1 [8960/60000 (15%)]\tLoss: 1.364717\n",
      "Train Epoch: 1 [9600/60000 (16%)]\tLoss: 1.344824\n",
      "Train Epoch: 1 [10240/60000 (17%)]\tLoss: 1.231810\n",
      "Train Epoch: 1 [10880/60000 (18%)]\tLoss: 1.212245\n",
      "Train Epoch: 1 [11520/60000 (19%)]\tLoss: 1.075095\n",
      "Train Epoch: 1 [12160/60000 (20%)]\tLoss: 1.151920\n",
      "Train Epoch: 1 [12800/60000 (21%)]\tLoss: 1.167744\n",
      "Train Epoch: 1 [13440/60000 (22%)]\tLoss: 1.074658\n",
      "Train Epoch: 1 [14080/60000 (23%)]\tLoss: 0.908480\n",
      "Train Epoch: 1 [14720/60000 (25%)]\tLoss: 0.839297\n",
      "Train Epoch: 1 [15360/60000 (26%)]\tLoss: 0.946202\n",
      "Train Epoch: 1 [16000/60000 (27%)]\tLoss: 0.951240\n",
      "Train Epoch: 1 [16640/60000 (28%)]\tLoss: 0.913577\n",
      "Train Epoch: 1 [17280/60000 (29%)]\tLoss: 0.963891\n",
      "Train Epoch: 1 [17920/60000 (30%)]\tLoss: 0.817639\n",
      "Train Epoch: 1 [18560/60000 (31%)]\tLoss: 0.757200\n",
      "Train Epoch: 1 [19200/60000 (32%)]\tLoss: 0.967922\n",
      "Train Epoch: 1 [19840/60000 (33%)]\tLoss: 0.854083\n",
      "Train Epoch: 1 [20480/60000 (34%)]\tLoss: 0.820100\n",
      "Train Epoch: 1 [21120/60000 (35%)]\tLoss: 0.658264\n",
      "Train Epoch: 1 [21760/60000 (36%)]\tLoss: 0.695334\n",
      "Train Epoch: 1 [22400/60000 (37%)]\tLoss: 0.601527\n",
      "Train Epoch: 1 [23040/60000 (38%)]\tLoss: 0.701636\n",
      "Train Epoch: 1 [23680/60000 (39%)]\tLoss: 0.639315\n",
      "Train Epoch: 1 [24320/60000 (41%)]\tLoss: 0.646477\n",
      "Train Epoch: 1 [24960/60000 (42%)]\tLoss: 0.668072\n",
      "Train Epoch: 1 [25600/60000 (43%)]\tLoss: 0.710025\n",
      "Train Epoch: 1 [26240/60000 (44%)]\tLoss: 0.484509\n",
      "Train Epoch: 1 [26880/60000 (45%)]\tLoss: 0.564650\n",
      "Train Epoch: 1 [27520/60000 (46%)]\tLoss: 0.603477\n",
      "Train Epoch: 1 [28160/60000 (47%)]\tLoss: 0.628251\n",
      "Train Epoch: 1 [28800/60000 (48%)]\tLoss: 0.644999\n",
      "Train Epoch: 1 [29440/60000 (49%)]\tLoss: 0.663249\n",
      "Train Epoch: 1 [30080/60000 (50%)]\tLoss: 0.527526\n",
      "Train Epoch: 1 [30720/60000 (51%)]\tLoss: 0.527981\n",
      "Train Epoch: 1 [31360/60000 (52%)]\tLoss: 0.577599\n",
      "Train Epoch: 1 [32000/60000 (53%)]\tLoss: 0.465394\n",
      "Train Epoch: 1 [32640/60000 (54%)]\tLoss: 0.471184\n",
      "Train Epoch: 1 [33280/60000 (55%)]\tLoss: 0.583723\n",
      "Train Epoch: 1 [33920/60000 (57%)]\tLoss: 0.558125\n",
      "Train Epoch: 1 [34560/60000 (58%)]\tLoss: 0.532155\n",
      "Train Epoch: 1 [35200/60000 (59%)]\tLoss: 0.470749\n",
      "Train Epoch: 1 [35840/60000 (60%)]\tLoss: 0.544474\n",
      "Train Epoch: 1 [36480/60000 (61%)]\tLoss: 0.514521\n",
      "Train Epoch: 1 [37120/60000 (62%)]\tLoss: 0.384303\n",
      "Train Epoch: 1 [37760/60000 (63%)]\tLoss: 0.444016\n",
      "Train Epoch: 1 [38400/60000 (64%)]\tLoss: 0.378163\n",
      "Train Epoch: 1 [39040/60000 (65%)]\tLoss: 0.382507\n",
      "Train Epoch: 1 [39680/60000 (66%)]\tLoss: 0.553740\n",
      "Train Epoch: 1 [40320/60000 (67%)]\tLoss: 0.343491\n",
      "Train Epoch: 1 [40960/60000 (68%)]\tLoss: 0.448462\n",
      "Train Epoch: 1 [41600/60000 (69%)]\tLoss: 0.476503\n",
      "Train Epoch: 1 [42240/60000 (70%)]\tLoss: 0.412186\n",
      "Train Epoch: 1 [42880/60000 (71%)]\tLoss: 0.385573\n",
      "Train Epoch: 1 [43520/60000 (72%)]\tLoss: 0.297511\n",
      "Train Epoch: 1 [44160/60000 (74%)]\tLoss: 0.455813\n",
      "Train Epoch: 1 [44800/60000 (75%)]\tLoss: 0.330928\n",
      "Train Epoch: 1 [45440/60000 (76%)]\tLoss: 0.355233\n",
      "Train Epoch: 1 [46080/60000 (77%)]\tLoss: 0.317835\n",
      "Train Epoch: 1 [46720/60000 (78%)]\tLoss: 0.512741\n",
      "Train Epoch: 1 [47360/60000 (79%)]\tLoss: 0.325811\n",
      "Train Epoch: 1 [48000/60000 (80%)]\tLoss: 0.529060\n",
      "Train Epoch: 1 [48640/60000 (81%)]\tLoss: 0.332176\n",
      "Train Epoch: 1 [49280/60000 (82%)]\tLoss: 0.367703\n",
      "Train Epoch: 1 [49920/60000 (83%)]\tLoss: 0.468644\n",
      "Train Epoch: 1 [50560/60000 (84%)]\tLoss: 0.424068\n",
      "Train Epoch: 1 [51200/60000 (85%)]\tLoss: 0.425136\n",
      "Train Epoch: 1 [51840/60000 (86%)]\tLoss: 0.342903\n",
      "Train Epoch: 1 [52480/60000 (87%)]\tLoss: 0.382762\n",
      "Train Epoch: 1 [53120/60000 (88%)]\tLoss: 0.506847\n",
      "Train Epoch: 1 [53760/60000 (90%)]\tLoss: 0.418062\n",
      "Train Epoch: 1 [54400/60000 (91%)]\tLoss: 0.428549\n",
      "Train Epoch: 1 [55040/60000 (92%)]\tLoss: 0.274689\n",
      "Train Epoch: 1 [55680/60000 (93%)]\tLoss: 0.388511\n",
      "Train Epoch: 1 [56320/60000 (94%)]\tLoss: 0.227903\n",
      "Train Epoch: 1 [56960/60000 (95%)]\tLoss: 0.448019\n",
      "Train Epoch: 1 [57600/60000 (96%)]\tLoss: 0.254313\n",
      "Train Epoch: 1 [58240/60000 (97%)]\tLoss: 0.232855\n",
      "Train Epoch: 1 [58880/60000 (98%)]\tLoss: 0.270845\n",
      "Train Epoch: 1 [59520/60000 (99%)]\tLoss: 0.311461\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\crsum\\Anaconda3\\envs\\pytorch\\lib\\site-packages\\torch\\nn\\_reduction.py:43: UserWarning: size_average and reduce args will be deprecated, please use reduction='sum' instead.\n",
      "  warnings.warn(warning.format(ret))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Test set: Avg. loss: 0.0199, Accuracy: 955/10000 (10%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0197, Accuracy: 1911/10000 (19%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0171, Accuracy: 2864/10000 (29%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0226, Accuracy: 3803/10000 (38%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0191, Accuracy: 4760/10000 (48%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0205, Accuracy: 5709/10000 (57%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0173, Accuracy: 6669/10000 (67%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0178, Accuracy: 7624/10000 (76%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0185, Accuracy: 8583/10000 (86%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0202, Accuracy: 9533/10000 (95%)\n",
      "\n",
      "Train Epoch: 2 [0/60000 (0%)]\tLoss: 0.270019\n",
      "Train Epoch: 2 [640/60000 (1%)]\tLoss: 0.452026\n",
      "Train Epoch: 2 [1280/60000 (2%)]\tLoss: 0.345250\n",
      "Train Epoch: 2 [1920/60000 (3%)]\tLoss: 0.352904\n",
      "Train Epoch: 2 [2560/60000 (4%)]\tLoss: 0.242325\n",
      "Train Epoch: 2 [3200/60000 (5%)]\tLoss: 0.324380\n",
      "Train Epoch: 2 [3840/60000 (6%)]\tLoss: 0.387620\n",
      "Train Epoch: 2 [4480/60000 (7%)]\tLoss: 0.380462\n",
      "Train Epoch: 2 [5120/60000 (9%)]\tLoss: 0.385470\n",
      "Train Epoch: 2 [5760/60000 (10%)]\tLoss: 0.306739\n",
      "Train Epoch: 2 [6400/60000 (11%)]\tLoss: 0.204381\n",
      "Train Epoch: 2 [7040/60000 (12%)]\tLoss: 0.382293\n",
      "Train Epoch: 2 [7680/60000 (13%)]\tLoss: 0.472363\n",
      "Train Epoch: 2 [8320/60000 (14%)]\tLoss: 0.427105\n",
      "Train Epoch: 2 [8960/60000 (15%)]\tLoss: 0.287996\n",
      "Train Epoch: 2 [9600/60000 (16%)]\tLoss: 0.248497\n",
      "Train Epoch: 2 [10240/60000 (17%)]\tLoss: 0.344463\n",
      "Train Epoch: 2 [10880/60000 (18%)]\tLoss: 0.206990\n",
      "Train Epoch: 2 [11520/60000 (19%)]\tLoss: 0.359905\n",
      "Train Epoch: 2 [12160/60000 (20%)]\tLoss: 0.275499\n",
      "Train Epoch: 2 [12800/60000 (21%)]\tLoss: 0.369505\n",
      "Train Epoch: 2 [13440/60000 (22%)]\tLoss: 0.251655\n",
      "Train Epoch: 2 [14080/60000 (23%)]\tLoss: 0.307185\n",
      "Train Epoch: 2 [14720/60000 (25%)]\tLoss: 0.242864\n",
      "Train Epoch: 2 [15360/60000 (26%)]\tLoss: 0.276642\n",
      "Train Epoch: 2 [16000/60000 (27%)]\tLoss: 0.310187\n",
      "Train Epoch: 2 [16640/60000 (28%)]\tLoss: 0.232993\n",
      "Train Epoch: 2 [17280/60000 (29%)]\tLoss: 0.324123\n",
      "Train Epoch: 2 [17920/60000 (30%)]\tLoss: 0.319054\n",
      "Train Epoch: 2 [18560/60000 (31%)]\tLoss: 0.240691\n",
      "Train Epoch: 2 [19200/60000 (32%)]\tLoss: 0.240464\n",
      "Train Epoch: 2 [19840/60000 (33%)]\tLoss: 0.227574\n",
      "Train Epoch: 2 [20480/60000 (34%)]\tLoss: 0.232879\n",
      "Train Epoch: 2 [21120/60000 (35%)]\tLoss: 0.152851\n",
      "Train Epoch: 2 [21760/60000 (36%)]\tLoss: 0.306894\n",
      "Train Epoch: 2 [22400/60000 (37%)]\tLoss: 0.209017\n",
      "Train Epoch: 2 [23040/60000 (38%)]\tLoss: 0.301205\n",
      "Train Epoch: 2 [23680/60000 (39%)]\tLoss: 0.190370\n",
      "Train Epoch: 2 [24320/60000 (41%)]\tLoss: 0.303294\n",
      "Train Epoch: 2 [24960/60000 (42%)]\tLoss: 0.368778\n",
      "Train Epoch: 2 [25600/60000 (43%)]\tLoss: 0.277961\n",
      "Train Epoch: 2 [26240/60000 (44%)]\tLoss: 0.329450\n",
      "Train Epoch: 2 [26880/60000 (45%)]\tLoss: 0.301529\n",
      "Train Epoch: 2 [27520/60000 (46%)]\tLoss: 0.406125\n",
      "Train Epoch: 2 [28160/60000 (47%)]\tLoss: 0.225504\n",
      "Train Epoch: 2 [28800/60000 (48%)]\tLoss: 0.248956\n",
      "Train Epoch: 2 [29440/60000 (49%)]\tLoss: 0.188870\n",
      "Train Epoch: 2 [30080/60000 (50%)]\tLoss: 0.410304\n",
      "Train Epoch: 2 [30720/60000 (51%)]\tLoss: 0.361421\n",
      "Train Epoch: 2 [31360/60000 (52%)]\tLoss: 0.331062\n",
      "Train Epoch: 2 [32000/60000 (53%)]\tLoss: 0.255686\n",
      "Train Epoch: 2 [32640/60000 (54%)]\tLoss: 0.343334\n",
      "Train Epoch: 2 [33280/60000 (55%)]\tLoss: 0.168918\n",
      "Train Epoch: 2 [33920/60000 (57%)]\tLoss: 0.223458\n",
      "Train Epoch: 2 [34560/60000 (58%)]\tLoss: 0.261126\n",
      "Train Epoch: 2 [35200/60000 (59%)]\tLoss: 0.280707\n",
      "Train Epoch: 2 [35840/60000 (60%)]\tLoss: 0.340354\n",
      "Train Epoch: 2 [36480/60000 (61%)]\tLoss: 0.286104\n",
      "Train Epoch: 2 [37120/60000 (62%)]\tLoss: 0.363793\n",
      "Train Epoch: 2 [37760/60000 (63%)]\tLoss: 0.284011\n",
      "Train Epoch: 2 [38400/60000 (64%)]\tLoss: 0.147259\n",
      "Train Epoch: 2 [39040/60000 (65%)]\tLoss: 0.249294\n",
      "Train Epoch: 2 [39680/60000 (66%)]\tLoss: 0.302720\n",
      "Train Epoch: 2 [40320/60000 (67%)]\tLoss: 0.282078\n",
      "Train Epoch: 2 [40960/60000 (68%)]\tLoss: 0.291829\n",
      "Train Epoch: 2 [41600/60000 (69%)]\tLoss: 0.272795\n",
      "Train Epoch: 2 [42240/60000 (70%)]\tLoss: 0.247200\n",
      "Train Epoch: 2 [42880/60000 (71%)]\tLoss: 0.475972\n",
      "Train Epoch: 2 [43520/60000 (72%)]\tLoss: 0.445721\n",
      "Train Epoch: 2 [44160/60000 (74%)]\tLoss: 0.261687\n",
      "Train Epoch: 2 [44800/60000 (75%)]\tLoss: 0.083440\n",
      "Train Epoch: 2 [45440/60000 (76%)]\tLoss: 0.203095\n",
      "Train Epoch: 2 [46080/60000 (77%)]\tLoss: 0.222996\n",
      "Train Epoch: 2 [46720/60000 (78%)]\tLoss: 0.258672\n",
      "Train Epoch: 2 [47360/60000 (79%)]\tLoss: 0.159625\n",
      "Train Epoch: 2 [48000/60000 (80%)]\tLoss: 0.284055\n",
      "Train Epoch: 2 [48640/60000 (81%)]\tLoss: 0.240361\n",
      "Train Epoch: 2 [49280/60000 (82%)]\tLoss: 0.217295\n",
      "Train Epoch: 2 [49920/60000 (83%)]\tLoss: 0.199004\n",
      "Train Epoch: 2 [50560/60000 (84%)]\tLoss: 0.304413\n",
      "Train Epoch: 2 [51200/60000 (85%)]\tLoss: 0.354952\n",
      "Train Epoch: 2 [51840/60000 (86%)]\tLoss: 0.213273\n",
      "Train Epoch: 2 [52480/60000 (87%)]\tLoss: 0.224339\n",
      "Train Epoch: 2 [53120/60000 (88%)]\tLoss: 0.354655\n",
      "Train Epoch: 2 [53760/60000 (90%)]\tLoss: 0.365776\n",
      "Train Epoch: 2 [54400/60000 (91%)]\tLoss: 0.255835\n",
      "Train Epoch: 2 [55040/60000 (92%)]\tLoss: 0.179386\n",
      "Train Epoch: 2 [55680/60000 (93%)]\tLoss: 0.213422\n",
      "Train Epoch: 2 [56320/60000 (94%)]\tLoss: 0.314783\n",
      "Train Epoch: 2 [56960/60000 (95%)]\tLoss: 0.173534\n",
      "Train Epoch: 2 [57600/60000 (96%)]\tLoss: 0.477465\n",
      "Train Epoch: 2 [58240/60000 (97%)]\tLoss: 0.184857\n",
      "Train Epoch: 2 [58880/60000 (98%)]\tLoss: 0.201993\n",
      "Train Epoch: 2 [59520/60000 (99%)]\tLoss: 0.190304\n",
      "\n",
      "Test set: Avg. loss: 0.0098, Accuracy: 973/10000 (10%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0115, Accuracy: 1942/10000 (19%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0103, Accuracy: 2911/10000 (29%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0091, Accuracy: 3890/10000 (39%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0098, Accuracy: 4866/10000 (49%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0127, Accuracy: 5826/10000 (58%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0110, Accuracy: 6794/10000 (68%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0094, Accuracy: 7766/10000 (78%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0112, Accuracy: 8730/10000 (87%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0105, Accuracy: 9700/10000 (97%)\n",
      "\n",
      "Train Epoch: 3 [0/60000 (0%)]\tLoss: 0.245565\n",
      "Train Epoch: 3 [640/60000 (1%)]\tLoss: 0.167261\n",
      "Train Epoch: 3 [1280/60000 (2%)]\tLoss: 0.226192\n",
      "Train Epoch: 3 [1920/60000 (3%)]\tLoss: 0.185219\n",
      "Train Epoch: 3 [2560/60000 (4%)]\tLoss: 0.384127\n",
      "Train Epoch: 3 [3200/60000 (5%)]\tLoss: 0.237584\n",
      "Train Epoch: 3 [3840/60000 (6%)]\tLoss: 0.272261\n",
      "Train Epoch: 3 [4480/60000 (7%)]\tLoss: 0.401726\n",
      "Train Epoch: 3 [5120/60000 (9%)]\tLoss: 0.193027\n",
      "Train Epoch: 3 [5760/60000 (10%)]\tLoss: 0.149906\n",
      "Train Epoch: 3 [6400/60000 (11%)]\tLoss: 0.196204\n",
      "Train Epoch: 3 [7040/60000 (12%)]\tLoss: 0.188998\n",
      "Train Epoch: 3 [7680/60000 (13%)]\tLoss: 0.141608\n",
      "Train Epoch: 3 [8320/60000 (14%)]\tLoss: 0.231801\n",
      "Train Epoch: 3 [8960/60000 (15%)]\tLoss: 0.305859\n",
      "Train Epoch: 3 [9600/60000 (16%)]\tLoss: 0.319822\n",
      "Train Epoch: 3 [10240/60000 (17%)]\tLoss: 0.240065\n",
      "Train Epoch: 3 [10880/60000 (18%)]\tLoss: 0.351428\n",
      "Train Epoch: 3 [11520/60000 (19%)]\tLoss: 0.175744\n",
      "Train Epoch: 3 [12160/60000 (20%)]\tLoss: 0.249372\n",
      "Train Epoch: 3 [12800/60000 (21%)]\tLoss: 0.100303\n",
      "Train Epoch: 3 [13440/60000 (22%)]\tLoss: 0.201921\n",
      "Train Epoch: 3 [14080/60000 (23%)]\tLoss: 0.135185\n",
      "Train Epoch: 3 [14720/60000 (25%)]\tLoss: 0.198703\n",
      "Train Epoch: 3 [15360/60000 (26%)]\tLoss: 0.169144\n",
      "Train Epoch: 3 [16000/60000 (27%)]\tLoss: 0.352172\n",
      "Train Epoch: 3 [16640/60000 (28%)]\tLoss: 0.208313\n",
      "Train Epoch: 3 [17280/60000 (29%)]\tLoss: 0.194136\n",
      "Train Epoch: 3 [17920/60000 (30%)]\tLoss: 0.360009\n",
      "Train Epoch: 3 [18560/60000 (31%)]\tLoss: 0.141690\n",
      "Train Epoch: 3 [19200/60000 (32%)]\tLoss: 0.256366\n",
      "Train Epoch: 3 [19840/60000 (33%)]\tLoss: 0.406252\n",
      "Train Epoch: 3 [20480/60000 (34%)]\tLoss: 0.140796\n",
      "Train Epoch: 3 [21120/60000 (35%)]\tLoss: 0.133917\n",
      "Train Epoch: 3 [21760/60000 (36%)]\tLoss: 0.173979\n",
      "Train Epoch: 3 [22400/60000 (37%)]\tLoss: 0.195618\n",
      "Train Epoch: 3 [23040/60000 (38%)]\tLoss: 0.130527\n",
      "Train Epoch: 3 [23680/60000 (39%)]\tLoss: 0.242225\n",
      "Train Epoch: 3 [24320/60000 (41%)]\tLoss: 0.283143\n",
      "Train Epoch: 3 [24960/60000 (42%)]\tLoss: 0.140506\n",
      "Train Epoch: 3 [25600/60000 (43%)]\tLoss: 0.259321\n",
      "Train Epoch: 3 [26240/60000 (44%)]\tLoss: 0.338801\n",
      "Train Epoch: 3 [26880/60000 (45%)]\tLoss: 0.172603\n",
      "Train Epoch: 3 [27520/60000 (46%)]\tLoss: 0.204689\n",
      "Train Epoch: 3 [28160/60000 (47%)]\tLoss: 0.263558\n",
      "Train Epoch: 3 [28800/60000 (48%)]\tLoss: 0.148636\n",
      "Train Epoch: 3 [29440/60000 (49%)]\tLoss: 0.173026\n",
      "Train Epoch: 3 [30080/60000 (50%)]\tLoss: 0.256834\n",
      "Train Epoch: 3 [30720/60000 (51%)]\tLoss: 0.380622\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 3 [31360/60000 (52%)]\tLoss: 0.231598\n",
      "Train Epoch: 3 [32000/60000 (53%)]\tLoss: 0.106008\n",
      "Train Epoch: 3 [32640/60000 (54%)]\tLoss: 0.200345\n",
      "Train Epoch: 3 [33280/60000 (55%)]\tLoss: 0.217677\n",
      "Train Epoch: 3 [33920/60000 (57%)]\tLoss: 0.090268\n",
      "Train Epoch: 3 [34560/60000 (58%)]\tLoss: 0.234017\n",
      "Train Epoch: 3 [35200/60000 (59%)]\tLoss: 0.196010\n",
      "Train Epoch: 3 [35840/60000 (60%)]\tLoss: 0.126041\n",
      "Train Epoch: 3 [36480/60000 (61%)]\tLoss: 0.119188\n",
      "Train Epoch: 3 [37120/60000 (62%)]\tLoss: 0.240342\n",
      "Train Epoch: 3 [37760/60000 (63%)]\tLoss: 0.194373\n",
      "Train Epoch: 3 [38400/60000 (64%)]\tLoss: 0.149166\n",
      "Train Epoch: 3 [39040/60000 (65%)]\tLoss: 0.154965\n",
      "Train Epoch: 3 [39680/60000 (66%)]\tLoss: 0.418261\n",
      "Train Epoch: 3 [40320/60000 (67%)]\tLoss: 0.269169\n",
      "Train Epoch: 3 [40960/60000 (68%)]\tLoss: 0.163501\n",
      "Train Epoch: 3 [41600/60000 (69%)]\tLoss: 0.129065\n",
      "Train Epoch: 3 [42240/60000 (70%)]\tLoss: 0.215807\n",
      "Train Epoch: 3 [42880/60000 (71%)]\tLoss: 0.186316\n",
      "Train Epoch: 3 [43520/60000 (72%)]\tLoss: 0.132331\n",
      "Train Epoch: 3 [44160/60000 (74%)]\tLoss: 0.187620\n",
      "Train Epoch: 3 [44800/60000 (75%)]\tLoss: 0.244651\n",
      "Train Epoch: 3 [45440/60000 (76%)]\tLoss: 0.266511\n",
      "Train Epoch: 3 [46080/60000 (77%)]\tLoss: 0.240693\n",
      "Train Epoch: 3 [46720/60000 (78%)]\tLoss: 0.155406\n",
      "Train Epoch: 3 [47360/60000 (79%)]\tLoss: 0.120030\n",
      "Train Epoch: 3 [48000/60000 (80%)]\tLoss: 0.288936\n",
      "Train Epoch: 3 [48640/60000 (81%)]\tLoss: 0.354993\n",
      "Train Epoch: 3 [49280/60000 (82%)]\tLoss: 0.163197\n",
      "Train Epoch: 3 [49920/60000 (83%)]\tLoss: 0.190297\n",
      "Train Epoch: 3 [50560/60000 (84%)]\tLoss: 0.208868\n",
      "Train Epoch: 3 [51200/60000 (85%)]\tLoss: 0.142231\n",
      "Train Epoch: 3 [51840/60000 (86%)]\tLoss: 0.159749\n",
      "Train Epoch: 3 [52480/60000 (87%)]\tLoss: 0.096229\n",
      "Train Epoch: 3 [53120/60000 (88%)]\tLoss: 0.230364\n",
      "Train Epoch: 3 [53760/60000 (90%)]\tLoss: 0.222976\n",
      "Train Epoch: 3 [54400/60000 (91%)]\tLoss: 0.297445\n",
      "Train Epoch: 3 [55040/60000 (92%)]\tLoss: 0.366638\n",
      "Train Epoch: 3 [55680/60000 (93%)]\tLoss: 0.175030\n",
      "Train Epoch: 3 [56320/60000 (94%)]\tLoss: 0.113492\n",
      "Train Epoch: 3 [56960/60000 (95%)]\tLoss: 0.100756\n",
      "Train Epoch: 3 [57600/60000 (96%)]\tLoss: 0.156181\n",
      "Train Epoch: 3 [58240/60000 (97%)]\tLoss: 0.186834\n",
      "Train Epoch: 3 [58880/60000 (98%)]\tLoss: 0.161804\n",
      "Train Epoch: 3 [59520/60000 (99%)]\tLoss: 0.116415\n",
      "\n",
      "Test set: Avg. loss: 0.0074, Accuracy: 983/10000 (10%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0088, Accuracy: 1952/10000 (20%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0094, Accuracy: 2921/10000 (29%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0075, Accuracy: 3900/10000 (39%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0081, Accuracy: 4876/10000 (49%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0085, Accuracy: 5851/10000 (59%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0071, Accuracy: 6831/10000 (68%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0093, Accuracy: 7804/10000 (78%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0082, Accuracy: 8781/10000 (88%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0073, Accuracy: 9762/10000 (98%)\n",
      "\n",
      "Train Epoch: 4 [0/60000 (0%)]\tLoss: 0.100080\n",
      "Train Epoch: 4 [640/60000 (1%)]\tLoss: 0.232660\n",
      "Train Epoch: 4 [1280/60000 (2%)]\tLoss: 0.156739\n",
      "Train Epoch: 4 [1920/60000 (3%)]\tLoss: 0.134120\n",
      "Train Epoch: 4 [2560/60000 (4%)]\tLoss: 0.301612\n",
      "Train Epoch: 4 [3200/60000 (5%)]\tLoss: 0.183098\n",
      "Train Epoch: 4 [3840/60000 (6%)]\tLoss: 0.128436\n",
      "Train Epoch: 4 [4480/60000 (7%)]\tLoss: 0.288487\n",
      "Train Epoch: 4 [5120/60000 (9%)]\tLoss: 0.224689\n",
      "Train Epoch: 4 [5760/60000 (10%)]\tLoss: 0.101360\n",
      "Train Epoch: 4 [6400/60000 (11%)]\tLoss: 0.172767\n",
      "Train Epoch: 4 [7040/60000 (12%)]\tLoss: 0.156306\n",
      "Train Epoch: 4 [7680/60000 (13%)]\tLoss: 0.375638\n",
      "Train Epoch: 4 [8320/60000 (14%)]\tLoss: 0.273568\n",
      "Train Epoch: 4 [8960/60000 (15%)]\tLoss: 0.197710\n",
      "Train Epoch: 4 [9600/60000 (16%)]\tLoss: 0.149683\n",
      "Train Epoch: 4 [10240/60000 (17%)]\tLoss: 0.247991\n",
      "Train Epoch: 4 [10880/60000 (18%)]\tLoss: 0.158113\n",
      "Train Epoch: 4 [11520/60000 (19%)]\tLoss: 0.195174\n",
      "Train Epoch: 4 [12160/60000 (20%)]\tLoss: 0.065742\n",
      "Train Epoch: 4 [12800/60000 (21%)]\tLoss: 0.175339\n",
      "Train Epoch: 4 [13440/60000 (22%)]\tLoss: 0.161262\n",
      "Train Epoch: 4 [14080/60000 (23%)]\tLoss: 0.226486\n",
      "Train Epoch: 4 [14720/60000 (25%)]\tLoss: 0.222933\n",
      "Train Epoch: 4 [15360/60000 (26%)]\tLoss: 0.119106\n",
      "Train Epoch: 4 [16000/60000 (27%)]\tLoss: 0.255322\n",
      "Train Epoch: 4 [16640/60000 (28%)]\tLoss: 0.217980\n",
      "Train Epoch: 4 [17280/60000 (29%)]\tLoss: 0.241015\n",
      "Train Epoch: 4 [17920/60000 (30%)]\tLoss: 0.143216\n",
      "Train Epoch: 4 [18560/60000 (31%)]\tLoss: 0.150457\n",
      "Train Epoch: 4 [19200/60000 (32%)]\tLoss: 0.090579\n",
      "Train Epoch: 4 [19840/60000 (33%)]\tLoss: 0.174999\n",
      "Train Epoch: 4 [20480/60000 (34%)]\tLoss: 0.273547\n",
      "Train Epoch: 4 [21120/60000 (35%)]\tLoss: 0.229463\n",
      "Train Epoch: 4 [21760/60000 (36%)]\tLoss: 0.113507\n",
      "Train Epoch: 4 [22400/60000 (37%)]\tLoss: 0.086400\n",
      "Train Epoch: 4 [23040/60000 (38%)]\tLoss: 0.115903\n",
      "Train Epoch: 4 [23680/60000 (39%)]\tLoss: 0.170052\n",
      "Train Epoch: 4 [24320/60000 (41%)]\tLoss: 0.104686\n",
      "Train Epoch: 4 [24960/60000 (42%)]\tLoss: 0.215830\n",
      "Train Epoch: 4 [25600/60000 (43%)]\tLoss: 0.171369\n",
      "Train Epoch: 4 [26240/60000 (44%)]\tLoss: 0.167831\n",
      "Train Epoch: 4 [26880/60000 (45%)]\tLoss: 0.154973\n",
      "Train Epoch: 4 [27520/60000 (46%)]\tLoss: 0.190393\n",
      "Train Epoch: 4 [28160/60000 (47%)]\tLoss: 0.144055\n",
      "Train Epoch: 4 [28800/60000 (48%)]\tLoss: 0.168517\n",
      "Train Epoch: 4 [29440/60000 (49%)]\tLoss: 0.300274\n",
      "Train Epoch: 4 [30080/60000 (50%)]\tLoss: 0.259732\n",
      "Train Epoch: 4 [30720/60000 (51%)]\tLoss: 0.130005\n",
      "Train Epoch: 4 [31360/60000 (52%)]\tLoss: 0.121172\n",
      "Train Epoch: 4 [32000/60000 (53%)]\tLoss: 0.146332\n",
      "Train Epoch: 4 [32640/60000 (54%)]\tLoss: 0.205992\n",
      "Train Epoch: 4 [33280/60000 (55%)]\tLoss: 0.253960\n",
      "Train Epoch: 4 [33920/60000 (57%)]\tLoss: 0.316268\n",
      "Train Epoch: 4 [34560/60000 (58%)]\tLoss: 0.258799\n",
      "Train Epoch: 4 [35200/60000 (59%)]\tLoss: 0.143042\n",
      "Train Epoch: 4 [35840/60000 (60%)]\tLoss: 0.084259\n",
      "Train Epoch: 4 [36480/60000 (61%)]\tLoss: 0.123645\n",
      "Train Epoch: 4 [37120/60000 (62%)]\tLoss: 0.109129\n",
      "Train Epoch: 4 [37760/60000 (63%)]\tLoss: 0.073098\n",
      "Train Epoch: 4 [38400/60000 (64%)]\tLoss: 0.141188\n",
      "Train Epoch: 4 [39040/60000 (65%)]\tLoss: 0.175083\n",
      "Train Epoch: 4 [39680/60000 (66%)]\tLoss: 0.200628\n",
      "Train Epoch: 4 [40320/60000 (67%)]\tLoss: 0.140696\n",
      "Train Epoch: 4 [40960/60000 (68%)]\tLoss: 0.212465\n",
      "Train Epoch: 4 [41600/60000 (69%)]\tLoss: 0.175765\n",
      "Train Epoch: 4 [42240/60000 (70%)]\tLoss: 0.204528\n",
      "Train Epoch: 4 [42880/60000 (71%)]\tLoss: 0.167871\n",
      "Train Epoch: 4 [43520/60000 (72%)]\tLoss: 0.250981\n",
      "Train Epoch: 4 [44160/60000 (74%)]\tLoss: 0.105432\n",
      "Train Epoch: 4 [44800/60000 (75%)]\tLoss: 0.141136\n",
      "Train Epoch: 4 [45440/60000 (76%)]\tLoss: 0.180265\n",
      "Train Epoch: 4 [46080/60000 (77%)]\tLoss: 0.173661\n",
      "Train Epoch: 4 [46720/60000 (78%)]\tLoss: 0.158791\n",
      "Train Epoch: 4 [47360/60000 (79%)]\tLoss: 0.134507\n",
      "Train Epoch: 4 [48000/60000 (80%)]\tLoss: 0.299999\n",
      "Train Epoch: 4 [48640/60000 (81%)]\tLoss: 0.190432\n",
      "Train Epoch: 4 [49280/60000 (82%)]\tLoss: 0.182624\n",
      "Train Epoch: 4 [49920/60000 (83%)]\tLoss: 0.395037\n",
      "Train Epoch: 4 [50560/60000 (84%)]\tLoss: 0.288223\n",
      "Train Epoch: 4 [51200/60000 (85%)]\tLoss: 0.244665\n",
      "Train Epoch: 4 [51840/60000 (86%)]\tLoss: 0.232603\n",
      "Train Epoch: 4 [52480/60000 (87%)]\tLoss: 0.224613\n",
      "Train Epoch: 4 [53120/60000 (88%)]\tLoss: 0.212172\n",
      "Train Epoch: 4 [53760/60000 (90%)]\tLoss: 0.155239\n",
      "Train Epoch: 4 [54400/60000 (91%)]\tLoss: 0.195784\n",
      "Train Epoch: 4 [55040/60000 (92%)]\tLoss: 0.071456\n",
      "Train Epoch: 4 [55680/60000 (93%)]\tLoss: 0.129284\n",
      "Train Epoch: 4 [56320/60000 (94%)]\tLoss: 0.168788\n",
      "Train Epoch: 4 [56960/60000 (95%)]\tLoss: 0.157610\n",
      "Train Epoch: 4 [57600/60000 (96%)]\tLoss: 0.115578\n",
      "Train Epoch: 4 [58240/60000 (97%)]\tLoss: 0.129425\n",
      "Train Epoch: 4 [58880/60000 (98%)]\tLoss: 0.205428\n",
      "Train Epoch: 4 [59520/60000 (99%)]\tLoss: 0.108520\n",
      "\n",
      "Test set: Avg. loss: 0.0077, Accuracy: 979/10000 (10%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0062, Accuracy: 1961/10000 (20%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0062, Accuracy: 2940/10000 (29%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0064, Accuracy: 3918/10000 (39%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0075, Accuracy: 4890/10000 (49%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0085, Accuracy: 5866/10000 (59%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0044, Accuracy: 6855/10000 (69%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0075, Accuracy: 7832/10000 (78%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0071, Accuracy: 8811/10000 (88%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0060, Accuracy: 9796/10000 (98%)\n",
      "\n",
      "Train Epoch: 5 [0/60000 (0%)]\tLoss: 0.246662\n",
      "Train Epoch: 5 [640/60000 (1%)]\tLoss: 0.134097\n",
      "Train Epoch: 5 [1280/60000 (2%)]\tLoss: 0.182673\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 5 [1920/60000 (3%)]\tLoss: 0.240178\n",
      "Train Epoch: 5 [2560/60000 (4%)]\tLoss: 0.227432\n",
      "Train Epoch: 5 [3200/60000 (5%)]\tLoss: 0.173642\n",
      "Train Epoch: 5 [3840/60000 (6%)]\tLoss: 0.180944\n",
      "Train Epoch: 5 [4480/60000 (7%)]\tLoss: 0.069482\n",
      "Train Epoch: 5 [5120/60000 (9%)]\tLoss: 0.114001\n",
      "Train Epoch: 5 [5760/60000 (10%)]\tLoss: 0.190454\n",
      "Train Epoch: 5 [6400/60000 (11%)]\tLoss: 0.212168\n",
      "Train Epoch: 5 [7040/60000 (12%)]\tLoss: 0.082159\n",
      "Train Epoch: 5 [7680/60000 (13%)]\tLoss: 0.129938\n",
      "Train Epoch: 5 [8320/60000 (14%)]\tLoss: 0.132226\n",
      "Train Epoch: 5 [8960/60000 (15%)]\tLoss: 0.175339\n",
      "Train Epoch: 5 [9600/60000 (16%)]\tLoss: 0.179772\n",
      "Train Epoch: 5 [10240/60000 (17%)]\tLoss: 0.218103\n",
      "Train Epoch: 5 [10880/60000 (18%)]\tLoss: 0.067516\n",
      "Train Epoch: 5 [11520/60000 (19%)]\tLoss: 0.102896\n",
      "Train Epoch: 5 [12160/60000 (20%)]\tLoss: 0.139666\n",
      "Train Epoch: 5 [12800/60000 (21%)]\tLoss: 0.108474\n",
      "Train Epoch: 5 [13440/60000 (22%)]\tLoss: 0.133580\n",
      "Train Epoch: 5 [14080/60000 (23%)]\tLoss: 0.092932\n",
      "Train Epoch: 5 [14720/60000 (25%)]\tLoss: 0.165435\n",
      "Train Epoch: 5 [15360/60000 (26%)]\tLoss: 0.060461\n",
      "Train Epoch: 5 [16000/60000 (27%)]\tLoss: 0.163205\n",
      "Train Epoch: 5 [16640/60000 (28%)]\tLoss: 0.204966\n",
      "Train Epoch: 5 [17280/60000 (29%)]\tLoss: 0.140494\n",
      "Train Epoch: 5 [17920/60000 (30%)]\tLoss: 0.131478\n",
      "Train Epoch: 5 [18560/60000 (31%)]\tLoss: 0.142646\n",
      "Train Epoch: 5 [19200/60000 (32%)]\tLoss: 0.088055\n",
      "Train Epoch: 5 [19840/60000 (33%)]\tLoss: 0.208320\n",
      "Train Epoch: 5 [20480/60000 (34%)]\tLoss: 0.262310\n",
      "Train Epoch: 5 [21120/60000 (35%)]\tLoss: 0.100460\n",
      "Train Epoch: 5 [21760/60000 (36%)]\tLoss: 0.076226\n",
      "Train Epoch: 5 [22400/60000 (37%)]\tLoss: 0.251690\n",
      "Train Epoch: 5 [23040/60000 (38%)]\tLoss: 0.239765\n",
      "Train Epoch: 5 [23680/60000 (39%)]\tLoss: 0.252496\n",
      "Train Epoch: 5 [24320/60000 (41%)]\tLoss: 0.150058\n",
      "Train Epoch: 5 [24960/60000 (42%)]\tLoss: 0.178571\n",
      "Train Epoch: 5 [25600/60000 (43%)]\tLoss: 0.135375\n",
      "Train Epoch: 5 [26240/60000 (44%)]\tLoss: 0.100756\n",
      "Train Epoch: 5 [26880/60000 (45%)]\tLoss: 0.074045\n",
      "Train Epoch: 5 [27520/60000 (46%)]\tLoss: 0.106955\n",
      "Train Epoch: 5 [28160/60000 (47%)]\tLoss: 0.216494\n",
      "Train Epoch: 5 [28800/60000 (48%)]\tLoss: 0.054652\n",
      "Train Epoch: 5 [29440/60000 (49%)]\tLoss: 0.212343\n",
      "Train Epoch: 5 [30080/60000 (50%)]\tLoss: 0.176648\n",
      "Train Epoch: 5 [30720/60000 (51%)]\tLoss: 0.147850\n",
      "Train Epoch: 5 [31360/60000 (52%)]\tLoss: 0.168981\n",
      "Train Epoch: 5 [32000/60000 (53%)]\tLoss: 0.139354\n",
      "Train Epoch: 5 [32640/60000 (54%)]\tLoss: 0.358169\n",
      "Train Epoch: 5 [33280/60000 (55%)]\tLoss: 0.149183\n",
      "Train Epoch: 5 [33920/60000 (57%)]\tLoss: 0.107161\n",
      "Train Epoch: 5 [34560/60000 (58%)]\tLoss: 0.166329\n",
      "Train Epoch: 5 [35200/60000 (59%)]\tLoss: 0.061547\n",
      "Train Epoch: 5 [35840/60000 (60%)]\tLoss: 0.161785\n",
      "Train Epoch: 5 [36480/60000 (61%)]\tLoss: 0.237391\n",
      "Train Epoch: 5 [37120/60000 (62%)]\tLoss: 0.233702\n",
      "Train Epoch: 5 [37760/60000 (63%)]\tLoss: 0.229163\n",
      "Train Epoch: 5 [38400/60000 (64%)]\tLoss: 0.144476\n",
      "Train Epoch: 5 [39040/60000 (65%)]\tLoss: 0.172338\n",
      "Train Epoch: 5 [39680/60000 (66%)]\tLoss: 0.155373\n",
      "Train Epoch: 5 [40320/60000 (67%)]\tLoss: 0.114263\n",
      "Train Epoch: 5 [40960/60000 (68%)]\tLoss: 0.097838\n",
      "Train Epoch: 5 [41600/60000 (69%)]\tLoss: 0.144523\n",
      "Train Epoch: 5 [42240/60000 (70%)]\tLoss: 0.314210\n",
      "Train Epoch: 5 [42880/60000 (71%)]\tLoss: 0.197592\n",
      "Train Epoch: 5 [43520/60000 (72%)]\tLoss: 0.341584\n",
      "Train Epoch: 5 [44160/60000 (74%)]\tLoss: 0.185514\n",
      "Train Epoch: 5 [44800/60000 (75%)]\tLoss: 0.207916\n",
      "Train Epoch: 5 [45440/60000 (76%)]\tLoss: 0.113827\n",
      "Train Epoch: 5 [46080/60000 (77%)]\tLoss: 0.195135\n",
      "Train Epoch: 5 [46720/60000 (78%)]\tLoss: 0.134169\n",
      "Train Epoch: 5 [47360/60000 (79%)]\tLoss: 0.245413\n",
      "Train Epoch: 5 [48000/60000 (80%)]\tLoss: 0.188531\n",
      "Train Epoch: 5 [48640/60000 (81%)]\tLoss: 0.298929\n",
      "Train Epoch: 5 [49280/60000 (82%)]\tLoss: 0.178754\n",
      "Train Epoch: 5 [49920/60000 (83%)]\tLoss: 0.082040\n",
      "Train Epoch: 5 [50560/60000 (84%)]\tLoss: 0.131702\n",
      "Train Epoch: 5 [51200/60000 (85%)]\tLoss: 0.232758\n",
      "Train Epoch: 5 [51840/60000 (86%)]\tLoss: 0.118410\n",
      "Train Epoch: 5 [52480/60000 (87%)]\tLoss: 0.135284\n",
      "Train Epoch: 5 [53120/60000 (88%)]\tLoss: 0.069190\n",
      "Train Epoch: 5 [53760/60000 (90%)]\tLoss: 0.138977\n",
      "Train Epoch: 5 [54400/60000 (91%)]\tLoss: 0.218364\n",
      "Train Epoch: 5 [55040/60000 (92%)]\tLoss: 0.083914\n",
      "Train Epoch: 5 [55680/60000 (93%)]\tLoss: 0.227464\n",
      "Train Epoch: 5 [56320/60000 (94%)]\tLoss: 0.099016\n",
      "Train Epoch: 5 [56960/60000 (95%)]\tLoss: 0.212852\n",
      "Train Epoch: 5 [57600/60000 (96%)]\tLoss: 0.276826\n",
      "Train Epoch: 5 [58240/60000 (97%)]\tLoss: 0.169828\n",
      "Train Epoch: 5 [58880/60000 (98%)]\tLoss: 0.215450\n",
      "Train Epoch: 5 [59520/60000 (99%)]\tLoss: 0.160372\n",
      "\n",
      "Test set: Avg. loss: 0.0072, Accuracy: 977/10000 (10%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0057, Accuracy: 1958/10000 (20%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0066, Accuracy: 2939/10000 (29%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0055, Accuracy: 3926/10000 (39%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0064, Accuracy: 4906/10000 (49%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0071, Accuracy: 5884/10000 (59%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0060, Accuracy: 6868/10000 (69%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0062, Accuracy: 7848/10000 (78%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0047, Accuracy: 8836/10000 (88%)\n",
      "\n",
      "\n",
      "Test set: Avg. loss: 0.0062, Accuracy: 9819/10000 (98%)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#Network training\n",
    "for epoch in range(1, n_epochs + 1):\n",
    "    train(epoch)\n",
    "    test()"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "## Loding the weights of the network trained above\n",
    "## This is useful for those who don't want to train the network afresh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Net(\n",
      "  (mnist): Sequential(\n",
      "    (0): Conv2d(1, 8, kernel_size=(3, 3), stride=(1, 1))\n",
      "    (1): ReLU()\n",
      "    (2): Conv2d(8, 16, kernel_size=(3, 3), stride=(1, 1))\n",
      "    (3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "    (5): ReLU()\n",
      "    (6): Dropout2d(p=0.5, inplace=False)\n",
      "    (7): Conv2d(16, 20, kernel_size=(5, 5), stride=(1, 1))\n",
      "    (8): BatchNorm2d(20, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "    (10): ReLU()\n",
      "    (11): Dropout2d(p=0.5, inplace=False)\n",
      "    (12): Flatten()\n",
      "    (13): Linear(in_features=320, out_features=50, bias=True)\n",
      "    (14): BatchNorm1d(50, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (15): ReLU()\n",
      "    (16): Linear(in_features=50, out_features=10, bias=True)\n",
      "    (17): LogSoftmax()\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "saved_weights = torch.load('model.pth')\n",
    "\n",
    "saved_model = Net()\n",
    "saved_model = saved_model\n",
    "saved_model = saved_model.cuda()\n",
    "\n",
    "saved_model.load_state_dict(saved_weights)\n",
    "saved_model.eval()\n",
    "print(saved_model)\n",
    "\n",
    "#input_image = cv2.imread('temp_folder/img_137.jpg',0)/255#example_data[2][0]\n",
    "#input_tensor = torch.tensor(input_image[np.newaxis,np.newaxis,:,:],dtype=torch.float32)\n",
    "#print(input_tensor)\n",
    "#output = saved_model(input_tensor.cuda())\n",
    "#print(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-1.3734e+01, -1.4056e+01, -1.2564e+01, -4.4823e-05, -1.4061e+01,\n",
      "         -1.0290e+01, -1.5320e+01, -1.3507e+01, -1.3215e+01, -1.3779e+01]],\n",
      "       device='cuda:0', grad_fn=<LogSoftmaxBackward>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\crsum\\Anaconda3\\envs\\pytorch\\lib\\site-packages\\torch\\nn\\modules\\container.py:92: UserWarning: Implicit dimension choice for log_softmax has been deprecated. Change the call to include dim=X as an argument.\n",
      "  input = module(input)\n"
     ]
    }
   ],
   "source": [
    "input_image = cv2.imread('temp_folder/img_137.jpg',0)/255#example_data[2][0]\n",
    "input_tensor = torch.tensor(input_image[np.newaxis,np.newaxis,:,:],dtype=torch.float32)\n",
    "#print(input_tensor)\n",
    "output = saved_model(input_tensor.cuda())\n",
    "print(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "conv_layer_idx = [0,2,7]\n",
    "relu_layer_idx = [1,5,10,15,16]\n",
    "fc_layer_idx = [13,16]\n",
    "batch_norm2d_idx = [3,8]\n",
    "batch_norm1d_idx = [14]\n",
    "max_pool_idx = [4,9]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge_batch_norm2d(conv_layer_idx, bn2d_idx):\n",
    "    eps =1e-05\n",
    "    conv_weight = saved_weights['mnist.'+str(conv_layer_idx)+'.weight']\n",
    "    conv_bias = saved_weights['mnist.'+str(conv_layer_idx)+'.bias']\n",
    "    bn2d_gamma = saved_weights['mnist.'+str(bn2d_idx)+'.weight']\n",
    "    bn2d_beta = saved_weights['mnist.'+str(bn2d_idx)+'.bias']\n",
    "    \n",
    "    bn2d_mean = saved_weights['mnist.'+str(bn2d_idx)+'.running_mean']\n",
    "    bn2d_var = saved_weights['mnist.'+str(bn2d_idx)+'.running_var']\n",
    "    bn2d_var = torch.sqrt(bn2d_var + eps)\n",
    "    print(conv_weight.shape)\n",
    "\n",
    "    conv_weight = conv_weight.transpose(0,3)\n",
    "    print(conv_weight.shape)\n",
    "    #gamma = bn.weight\n",
    "    #beta = bn.bias\n",
    "    #mean = bn.running_mean\n",
    "    #var = bn.running_var\n",
    "    #eps =1e-05\n",
    "\n",
    "    #var_sqrt = torch.sqrt(var + eps)\n",
    "\n",
    "    #w = (self.weight * gamma.reshape(self.out_channels, 1, 1, 1)) / var_sqrt.reshape(self.out_channels, 1,1, 1)\n",
    "    #b = ((self.bias - mean) * gamma) / var_sqrt + beta\n",
    "\n",
    "    output_conv_weight = torch.mul(conv_weight, torch.div(bn2d_gamma, bn2d_var))\n",
    "    #print(output_conv_weight.shape)\n",
    "    output_conv_weight = output_conv_weight.transpose(0,3)\n",
    "    conv_weight = conv_weight.transpose(0,3)\n",
    "    output_conv_bias = torch.add(torch.mul((conv_bias-bn2d_mean), torch.div(bn2d_gamma,bn2d_var)),bn2d_beta)\n",
    "    return output_conv_weight, output_conv_bias#, #output_conv_weight, \n",
    "\n",
    "def merge_batch_norm1d(fc_layer_idx, bn1d_idx):\n",
    "    eps =1e-05\n",
    "\n",
    "    fc_weight = saved_weights['mnist.'+str(fc_layer_idx)+'.weight']\n",
    "    fc_bias = saved_weights['mnist.'+str(fc_layer_idx)+'.bias']\n",
    "    bn1d_gamma = saved_weights['mnist.'+str(bn1d_idx)+'.weight']\n",
    "    bn1d_beta = saved_weights['mnist.'+str(bn1d_idx)+'.bias']\n",
    "    bn1d_mean = saved_weights['mnist.'+str(bn1d_idx)+'.running_mean']\n",
    "    bn1d_var = saved_weights['mnist.'+str(bn1d_idx)+'.running_var']\n",
    "    bn1d_var = torch.sqrt(bn1d_var + eps)\n",
    "    fc_weight = fc_weight.transpose(0,1)\n",
    "    output_fc_weight = torch.mul(fc_weight, torch.div(bn1d_gamma, bn1d_var))\n",
    "    fc_weight = fc_weight.transpose(0,1)\n",
    "    output_fc_weight = output_fc_weight.transpose(0,1)\n",
    "    output_fc_bias = torch.add(torch.mul((fc_bias-bn1d_mean), torch.div(bn1d_gamma,bn1d_var)),bn1d_beta)\n",
    "    return output_fc_weight, output_fc_bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def output_list_fn(layers_list,data_loader,model):\n",
    "    tensor_list = []\n",
    "    output_dict = {}\n",
    "    count =0\n",
    "    for batch_idx, (data, _) in enumerate(data_loader):\n",
    "        max_output_list =[]\n",
    "        x = data.cpu()\n",
    "        for idx,seq in enumerate(model):#.mnist):\n",
    "            x = seq(x)\n",
    "            if idx in layers_list:\n",
    "                max_output_list.append(x) #intermediate_output#(torch.flatten(x)))\n",
    "        #del data_cuda\n",
    "        #torch.cuda.empty_cache()\n",
    "        iter = 0\n",
    "        if count == 0:\n",
    "            tensor_list.append(data)\n",
    "        else:\n",
    "            tensor_list[iter] = torch.cat((tensor_list[iter],data),0)\n",
    "        iter+=1   \n",
    "        for idx in layers_list:\n",
    "            if count == 0:\n",
    "                tensor_list.append(max_output_list[iter-1])\n",
    "            else:\n",
    "                tensor_list[iter] = torch.cat((tensor_list[iter],max_output_list[iter-1]),0)\n",
    "            iter+=1\n",
    "        \n",
    "        count +=1\n",
    "    tensor_list[0] = torch.flatten(tensor_list[0]).detach().numpy()\n",
    "    output_dict['input'] = tensor_list[0]\n",
    "    iter =1\n",
    "    for idx in layers_list:\n",
    "        tensor_list[iter] = torch.flatten(tensor_list[iter]).detach().numpy()\n",
    "        output_dict[str(idx)] = tensor_list[iter]\n",
    "        iter+=1\n",
    "\n",
    "    return output_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def output_list_fn_cuda(layers_list,data_loader,model):\n",
    "    tensor_list = []\n",
    "    output_dict = {}\n",
    "    count =0\n",
    "    for batch_idx, (data, _) in enumerate(data_loader):\n",
    "        max_output_list =[]\n",
    "        x = data.cuda()\n",
    "        for idx,seq in enumerate(model):#.mnist):\n",
    "            x = seq(x)\n",
    "            if idx in layers_list:\n",
    "                max_output_list.append(x.detach().cpu()) #intermediate_output#(torch.flatten(x)))\n",
    "        #del data_cuda\n",
    "        #torch.cuda.empty_cache()\n",
    "        iter = 0\n",
    "        if count == 0:\n",
    "            tensor_list.append(data)\n",
    "        else:\n",
    "            tensor_list[iter] = torch.cat((tensor_list[iter],data.cpu()),0)\n",
    "        iter+=1   \n",
    "        for idx in layers_list:\n",
    "            if count == 0:\n",
    "                tensor_list.append(max_output_list[iter-1])\n",
    "            else:\n",
    "                tensor_list[iter] = torch.cat((tensor_list[iter],max_output_list[iter-1]),0)\n",
    "            iter+=1\n",
    "        \n",
    "        count +=1\n",
    "    tensor_list[0] = torch.flatten(tensor_list[0]).detach().numpy()\n",
    "    output_dict['input'] = tensor_list[0]\n",
    "    iter =1\n",
    "    for idx in layers_list:\n",
    "        tensor_list[iter] = torch.flatten(tensor_list[iter]).detach().numpy()\n",
    "        output_dict[str(idx)] = tensor_list[iter]\n",
    "        iter+=1\n",
    "    \n",
    "    \n",
    "    return output_dict\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\crsum\\Anaconda3\\envs\\pytorch\\lib\\site-packages\\ipykernel_launcher.py:9: UserWarning: Implicit dimension choice for log_softmax has been deprecated. Change the call to include dim=X as an argument.\n",
      "  if __name__ == '__main__':\n"
     ]
    }
   ],
   "source": [
    "relu_output_dict = output_list_fn_cuda(relu_layer_idx,train_loader,saved_model.mnist.cuda())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "99 percentile 0.9960784316062927\n",
      "99.9 percentile 1.0\n",
      "99.99 percentile 1.0\n",
      "100 percentile 1.0\n",
      "99 percentile 0.657135131955144\n",
      "99.9 percentile 0.816114008665096\n",
      "99.99 percentile 0.9123463034689401\n",
      "100 percentile 1.0883766412734985\n",
      "99 percentile 2.936004161834717\n",
      "99.9 percentile 3.583656788110737\n",
      "99.99 percentile 4.141691210078903\n",
      "100 percentile 6.034055233001709\n",
      "99 percentile 2.37951397895813\n",
      "99.9 percentile 2.841526761770279\n",
      "99.99 percentile 3.1277217894553537\n",
      "100 percentile 3.9231252670288086\n",
      "99 percentile 2.923907430171963\n",
      "99.9 percentile 3.579677352190019\n",
      "99.99 percentile 4.048984072303735\n",
      "100 percentile 4.745292663574219\n",
      "99 percentile 10.84741925239563\n",
      "99.9 percentile 12.236660329818786\n",
      "99.99 percentile 13.058166522598256\n",
      "100 percentile 14.453241348266602\n"
     ]
    }
   ],
   "source": [
    "# print percentiles\n",
    "relu_layers = [1,5,10,15,16]\n",
    "percentile_99 = np.percentile(relu_output_dict['input'],99)\n",
    "percentile_99_9 = np.percentile(relu_output_dict['input'],99.9)\n",
    "percentile_99_99 = np.percentile(relu_output_dict['input'],99.99)\n",
    "percentile_100 = np.percentile(relu_output_dict['input'],100)\n",
    "print('99 percentile', percentile_99)\n",
    "print('99.9 percentile', percentile_99_9)\n",
    "print('99.99 percentile', percentile_99_99)\n",
    "print('100 percentile', percentile_100)\n",
    "percentile_dict = {}\n",
    "percentile_dict['layer_input_99'] = percentile_99\n",
    "percentile_dict['layer_input_99_9'] = percentile_99_9\n",
    "percentile_dict['layer_input_99_99'] = percentile_99_99\n",
    "percentile_dict['layer_input_100'] = percentile_100\n",
    "for idx in relu_layers:\n",
    "    percentile_99 = np.percentile(relu_output_dict[str(idx)],99)\n",
    "    percentile_99_9 = np.percentile(relu_output_dict[str(idx)],99.9)\n",
    "    percentile_99_99 = np.percentile(relu_output_dict[str(idx)],99.99)\n",
    "    percentile_100 = np.percentile(relu_output_dict[str(idx)],100)\n",
    "    percentile_dict['layer_'+str(idx)+'_99'] = percentile_99\n",
    "    percentile_dict['layer_'+str(idx)+'_99_9'] = percentile_99_9\n",
    "    percentile_dict['layer_'+str(idx)+'_99_99'] = percentile_99_99\n",
    "    percentile_dict['layer_'+str(idx)+'_100'] = percentile_100\n",
    "\n",
    "    print('99 percentile', percentile_99)\n",
    "    print('99.9 percentile', percentile_99_9)\n",
    "    print('99.99 percentile', percentile_99_99)\n",
    "\n",
    "    print('100 percentile', percentile_100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize_weight_bias(weight,bias,l1_out,l2_out):\n",
    "    weight = torch.mul(weight,l1_out/l2_out)\n",
    "    bias = torch.div(bias,l2_out)\n",
    "    return weight,bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 8, 3, 3])\n",
      "torch.Size([3, 8, 3, 16])\n",
      "torch.Size([20, 16, 5, 5])\n",
      "torch.Size([5, 16, 5, 20])\n"
     ]
    }
   ],
   "source": [
    "conv1_weight = saved_weights['mnist.0.weight']\n",
    "conv1_bias = saved_weights['mnist.0.bias']\n",
    "\n",
    "conv2_weight, conv2_bias = merge_batch_norm2d(2,3)\n",
    "conv3_weight, conv3_bias = merge_batch_norm2d(7,8)\n",
    "fc1_weight,fc1_bias = merge_batch_norm1d(13,14)\n",
    "fc2_weight = saved_weights['mnist.16.weight']\n",
    "fc2_bias = saved_weights['mnist.16.bias']\n",
    "\n",
    "snn_conv1_weight,snn_conv1_bias = normalize_weight_bias(conv1_weight,conv1_bias,percentile_dict['layer_input_99_9'],percentile_dict['layer_1_99_9'])\n",
    "snn_conv2_weight,snn_conv2_bias = normalize_weight_bias(conv2_weight,conv2_bias,percentile_dict['layer_1_99_9'],percentile_dict['layer_5_99_9'])\n",
    "snn_conv3_weight,snn_conv3_bias = normalize_weight_bias(conv3_weight,conv3_bias,percentile_dict['layer_5_99_9'],percentile_dict['layer_10_99_9'])\n",
    "snn_fc1_weight,snn_fc1_bias = normalize_weight_bias(fc1_weight,fc1_bias,percentile_dict['layer_10_99_9'],percentile_dict['layer_15_99_9'])\n",
    "snn_fc2_weight,snn_fc2_bias = normalize_weight_bias(fc2_weight,fc2_bias,percentile_dict['layer_15_99_9'],percentile_dict['layer_16_99_9'])\n",
    "\n",
    "\n",
    "\n",
    "#print(snn_conv1_weight)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_layer_idx = [0,2,7]\n",
    "relu_layer_idx = [1,5,10,15,16]\n",
    "fc_layer_idx = [13,16]\n",
    "batch_norm2d_idx = [3,8]\n",
    "batch_norm1d_idx = [14]\n",
    "max_pool_idx = [4,9]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_batch_size = 2000\n",
    "test_batch_size = 1000\n",
    "snn_train_loader = torch.utils.data.DataLoader(\n",
    "  #torchvision.datasets.MNIST('/files/', train=True, download=True,\n",
    "  torchvision.datasets.MNIST('.', train=True, download=True,\n",
    "\n",
    "                             transform=torchvision.transforms.Compose([\n",
    "                               torchvision.transforms.ToTensor()\n",
    "                             ])),\n",
    "  batch_size=train_batch_size,pin_memory = False, shuffle=True)\n",
    "\n",
    "snn_test_loader = torch.utils.data.DataLoader(\n",
    "  #torchvision.datasets.MNIST('/files/', train=False, download=True,\n",
    "  torchvision.datasets.MNIST('.', train=False, download=True,\n",
    "\n",
    "                             transform=torchvision.transforms.Compose([\n",
    "                               torchvision.transforms.ToTensor()\n",
    "                             ])),\n",
    "  batch_size=test_batch_size, pin_memory = False,shuffle=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "60000\n",
      "10000\n"
     ]
    }
   ],
   "source": [
    "train_size =0\n",
    "for batch_idx, (data, target) in enumerate(snn_train_loader):\n",
    "    train_size += data.shape[0]\n",
    "print(train_size)\n",
    "\n",
    "test_size =0\n",
    "for batch_idx, (data, target) in enumerate(snn_test_loader):\n",
    "    test_size += data.shape[0]\n",
    "print(test_size)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Data gen: Generate data and analyze how many time steps on an average is needed for someone to confirm the output\n",
    "#spikenet_fc_1 = spiking_fc(1,snn_fc1_weight,snn_fc1_bias,1).cuda()\n",
    "#spikenet_fc_2 = spiking_fc(1,snn_fc2_weight,snn_fc2_bias,1).cuda()\n",
    "#spike_frame = torch.zeros_like(input_tensor)\n",
    "#spike_pot = torch.zeros_like(input_tensor)\n",
    "sp1_shape = snn_conv1_weight.shape\n",
    "spikenet_1 = nn.Conv2d(sp1_shape[0],sp1_shape[1],kernel_size=sp1_shape[2]).cuda()\n",
    "spikenet_1.weight = nn.Parameter(snn_conv1_weight)\n",
    "spikenet_1.bias = nn.Parameter(snn_conv1_bias)\n",
    "sp2_shape = snn_conv2_weight.shape\n",
    "spikenet_2 = nn.Conv2d(sp2_shape[0],sp2_shape[1],kernel_size=sp2_shape[2]).cuda()\n",
    "spikenet_2.weight = nn.Parameter(snn_conv2_weight)\n",
    "spikenet_2.bias = nn.Parameter(snn_conv2_bias)\n",
    "sp3_shape = snn_conv3_weight.shape\n",
    "spikenet_3 = nn.Conv2d(sp3_shape[0],sp3_shape[1],kernel_size=sp3_shape[2]).cuda()\n",
    "spikenet_3.weight = nn.Parameter(snn_conv3_weight)\n",
    "spikenet_3.bias = nn.Parameter(snn_conv3_bias)\n",
    "sp_fc1_shape = snn_fc1_weight.shape\n",
    "spikenet_fc_1 = nn.Linear(sp_fc1_shape[0],sp_fc1_shape[1]).cuda()\n",
    "spikenet_fc_1.weight = nn.Parameter(snn_fc1_weight)\n",
    "spikenet_fc_1.bias = nn.Parameter(snn_fc1_bias)\n",
    "sp_fc2_shape = snn_fc2_weight.shape\n",
    "spikenet_fc_2 = nn.Linear(sp_fc2_shape[0],sp_fc2_shape[1]).cuda()\n",
    "spikenet_fc_2.weight = nn.Parameter(snn_fc2_weight)\n",
    "spikenet_fc_2.bias = nn.Parameter(snn_fc2_bias)\n",
    "#max_pool2 = nn.MaxPool2d(2, stride=2).cuda()\n",
    "#max_pool3 = nn.MaxPool2d(2, stride=2).cuda()\n",
    "max_pool2 = nn.MaxPool2d(2,stride=2,return_indices=True).cuda()\n",
    "max_unpool2 = nn.MaxUnpool2d(2, stride=2).cuda()\n",
    "max_pool3 = nn.MaxPool2d(2,stride=2,return_indices=True).cuda()\n",
    "max_unpool3 = nn.MaxUnpool2d(2, stride=2).cuda()\n",
    "\n",
    "number_of_timesteps = 500 #1000\n",
    "train_output_expected = np.zeros((train_size,1))\n",
    "train_output_array = np.zeros((train_size,number_of_timesteps))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Running SNN on train data "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_timesteps = 500\n",
    "test_output_expected = np.zeros((test_size,1))\n",
    "test_output_array = np.zeros((test_size,number_of_timesteps))\n",
    "\n",
    "for batch_idx, (data, target) in enumerate(snn_train_loader):\n",
    "    #x = data np.zeros((1000,1000))\n",
    "    start = batch_idx*train_batch_size\n",
    "    end = batch_idx*train_batch_size + data.shape[0]\n",
    "    test_output_expected[start:end,0] = target.detach().numpy()\n",
    "    data = data.cuda()\n",
    "    #print(data.shape)\n",
    "    spike_pot = torch.zeros_like(data).detach()cuda()\n",
    "    sp1_out = torch.zeros((data.shape[0],8,26,26)).detach().cuda()\n",
    "    sp2_out = torch.zeros((data.shape[0],16,24,24)).detach().cuda()\n",
    "    sp3_out = torch.zeros((data.shape[0],20,8,8)).detach().cuda()\n",
    "    sp_fc1_out = torch.zeros((data.shape[0],50)).detach().cuda()\n",
    "    sp_fc2_out = torch.zeros((data.shape[0],10)).detach().cuda()\n",
    "\n",
    "    \n",
    "    for i in range(number_of_timesteps):\n",
    "        print(i)\n",
    "        first_frame = True if i ==0 else False\n",
    "        #x = spike_frame_mnist.reset_by_subtraction(input_tensor,first_frame)\n",
    "            \n",
    "        #spike_pot = torch.add(spike_pot, data)\n",
    "        spike_pot.add_(data)\n",
    "\n",
    "        spike_frame = torch.gt(spike_pot,1.0)*1.0\n",
    "        spike_pot.sub_(spike_frame)\n",
    "        sp1_out.add_(spikenet_1(spike_frame))\n",
    "        sp2_in = torch.gt(sp1_out,1.0)*1.0\n",
    "        sp1_out.sub_(sp2_in)\n",
    "        sp2_out.add_(spikenet_2(sp2_in))\n",
    "        sp3_in = torch.gt(sp2_out,1.0)*1.0\n",
    "        sp2_out.sub_(sp3_in)\n",
    "        sp3_in_red = max_pool2(sp3_in)\n",
    "        sp3_out.add_(spikenet_3(sp3_in_red))\n",
    "        sp4_in = torch.gt(sp3_out,1.0)*1.0\n",
    "        sp3_out.sub_(sp4_in)\n",
    "        sp4_in_red = max_pool3(sp4_in)\n",
    "        sp4_in_red_flat = torch.flatten(sp4_in_red,start_dim=1)\n",
    "        sp_fc1_out.add_(spikenet_fc_1(sp4_in_red_flat))\n",
    "        sp_fc2_in = torch.gt(sp_fc1_out,1.0)*1.0\n",
    "        sp_fc1_out.sub_(sp_fc2_in)\n",
    "        sp_fc2_out.add_(spikenet_fc_2(sp_fc2_in))\n",
    "        sof_pot = F.softmax(sp_fc2_out)\n",
    "        max_numbers = torch.argmax(sof_pot,dim=1)\n",
    "        test_output_array[start:end,i] = max_numbers.cpu().detach().numpy()\n",
    "        \n",
    "        torch.cuda.empty_cache()\n",
    "    torch.cuda.empty_cache()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_layer_for_snn_pool(num_inp_ch):\n",
    "    conv = nn.Conv2d(num_inp_ch,num_inp_ch,kernel_size=2,padding=0,stride=2).cuda()\n",
    "    weight = torch.zeros((num_inp_ch,num_inp_ch,2,2)).cuda()\n",
    "    for i in range(num_inp_ch):\n",
    "        weight[i,i,:,:].add_(1.0)\n",
    "    bias = torch.zeros(num_inp_ch).cuda()\n",
    "    conv.weight = nn.Parameter(weight)\n",
    "    conv.bias = nn.Parameter(bias)\n",
    "    return conv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cnn2gemm(input_frame,kernel_size,padding=True):\n",
    "    input_shape = input_frame.shape\n",
    "    output_shape = input_frame.shape\n",
    "    input_frame_tmp = None\n",
    "    padding = (kernel_size+1)//2\n",
    "    if not padding:\n",
    "        output_shape[2] -= padding\n",
    "        output_shape[3] -= padding\n",
    "    input_frame_tmp= np.zeros((*output_shape),dtype= np.float32)\n",
    "    \n",
    "    if padding:\n",
    "        input_frame_tmp[:,:,(padding//2):-(padding//2),(padding//2):-(padding//2)] = input_frame\n",
    "    output_array = np.zeros((output_shape[1]*kernel_size*kernel_size,output_shape[2]*output_shape[3]),dtype= np.float32)\n",
    "    \n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Running SNN on test data "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\crsum\\Anaconda3\\envs\\pytorch\\lib\\site-packages\\ipykernel_launcher.py:66: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "2\n",
      "3\n",
      "4\n",
      "5\n",
      "6\n",
      "7\n",
      "8\n",
      "9\n"
     ]
    }
   ],
   "source": [
    "### RAM PLEASE LOOK AT THIS CODE\n",
    "\n",
    "number_of_timesteps = 400\n",
    "sp2_pool_conv = conv_layer_for_snn_pool(16)\n",
    "sp3_pool_conv = conv_layer_for_snn_pool(20)\n",
    "\n",
    "test_output_expected = np.zeros((test_size,1))\n",
    "test_output_array = np.zeros((test_size,number_of_timesteps))\n",
    "for batch_idx, (data, target) in enumerate(snn_test_loader):\n",
    "    start = batch_idx*test_batch_size\n",
    "    end = batch_idx*test_batch_size + data.shape[0]\n",
    "    test_output_expected[start:end,0] = target.detach().numpy()\n",
    "    data = data.cuda()\n",
    "    spike_pot = torch.zeros_like(data).detach().cuda()\n",
    "    sp1_np = np.zeros((data.shape[0],8,26,26))\n",
    "    sp1_out = torch.tensor(sp1_np, requires_grad=False, dtype=torch.float32).cuda()\n",
    "    sp2_np = np.zeros((data.shape[0],16,24,24))\n",
    "    sp2_out = torch.tensor(sp2_np, requires_grad=False, dtype=torch.float32).cuda()\n",
    "    sp2_pool = torch.tensor(sp2_np, requires_grad=False, dtype=torch.float32).cuda()\n",
    "    sp2_unpool = torch.tensor(sp2_np, requires_grad=False, dtype=torch.float32).cuda()\n",
    "    sp3_np = np.zeros((data.shape[0],20,8,8))\n",
    "    sp3_out = torch.tensor(sp3_np, requires_grad=False, dtype=torch.float32).cuda()\n",
    "    sp3_pool = torch.tensor(sp3_np, requires_grad=False, dtype=torch.float32).cuda()\n",
    "    sp3_unpool = torch.tensor(sp3_np, requires_grad=False, dtype=torch.float32).cuda()\n",
    "    sp_fc1_np = np.zeros((data.shape[0],50))\n",
    "    sp_fc1_out = torch.tensor(sp_fc1_np, requires_grad=False, dtype=torch.float32).cuda()\n",
    "    sp_fc2_np = np.zeros((data.shape[0],10))\n",
    "    sp_fc2_out = torch.tensor(sp_fc2_np, requires_grad=False, dtype=torch.float32).cuda()\n",
    "    print(batch_idx)\n",
    "    for i in range(number_of_timesteps):\n",
    "        #print(i)\n",
    "        first_frame = True if i ==0 else False\n",
    "        spike_pot.add_(data)\n",
    "\n",
    "        spike_frame = torch.gt(spike_pot,1.0)*1.0 #F.threshold_(spike_pot,1,0).sign() #\n",
    "        spike_pot.sub_(spike_frame)\n",
    "        sp1_out.add_(spikenet_1(spike_frame).detach())\n",
    "        sp2_in = torch.gt(sp1_out,1.0)*1.0 #F.threshold_(sp1_out, 1, 0) #\n",
    "        sp1_out.sub_(sp2_in)\n",
    "        sp2_out.add_(spikenet_2(sp2_in).detach())\n",
    "        sp3_in = torch.gt(sp2_out,1.0)*1.0 #F.threshold_(sp2_out, 1, 0).sign() # \n",
    "        sp2_out.sub_(sp3_in)\n",
    "        sp2_pool.add_(sp3_in)\n",
    "        sp2_pool_out = max_pool2(sp2_pool)\n",
    "        sp2_unpool_out = max_unpool2(*sp2_pool_out)\n",
    "        sp2_pool_non_zero = (sp2_unpool_out != 0)*1.0\n",
    "        sp3_in_masked = torch.mul(sp3_in,sp2_pool_non_zero)\n",
    "        sp3_in_red = sp2_pool_conv(sp3_in_masked)\n",
    "        #sp3_in_red = max_pool2(sp3_in)\n",
    "        sp3_out.add_(spikenet_3(sp3_in_red).detach())\n",
    "        sp4_in = torch.gt(sp3_out,1.0)*1.0 #F.threshold_(sp3_out, 1, 0).sign() #\n",
    "        sp3_out.sub_(sp4_in)\n",
    "        sp3_pool.add_(sp4_in)\n",
    "        sp3_pool_out = max_pool3(sp3_pool)\n",
    "        sp3_unpool_out = max_unpool3(*sp3_pool_out)\n",
    "        sp3_pool_non_zero = (sp3_unpool_out != 0)*1.0\n",
    "        sp4_in_masked = torch.mul(sp4_in,sp3_pool_non_zero)\n",
    "        sp4_in_red = sp3_pool_conv(sp4_in_masked)\n",
    "        \n",
    "        #sp4_in_red = max_pool3(sp4_in)\n",
    "        sp4_in_red_flat = torch.flatten(sp4_in_red,start_dim=1)\n",
    "        sp_fc1_out.add_(spikenet_fc_1(sp4_in_red_flat).detach())\n",
    "        sp_fc2_in = torch.gt(sp_fc1_out,1.0)*1.0 #F.threshold_(sp_fc1_out, 1, 0).sign() #\n",
    "        sp_fc1_out.sub_(sp_fc2_in)\n",
    "        sp_fc2_out.add_(spikenet_fc_2(sp_fc2_in).detach())\n",
    "        sof_pot = F.softmax(sp_fc2_out).detach()\n",
    "        max_numbers = torch.argmax(sof_pot,dim=1)\n",
    "        test_output_array[start:end,i] = max_numbers.cpu().detach().numpy()\n",
    "        #print(torch.cuda.memory_allocated(device=0))\n",
    "        torch.cuda.empty_cache()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "accuracy_over_time = np.zeros(test_output_array.shape[1])\n",
    "total_samples = test_output_array.shape[0]\n",
    "from scipy import stats\n",
    "for idx in range(test_output_array.shape[1]):\n",
    "    part_of_output = test_output_array[:,0:idx+1]\n",
    "    max_of_part = stats.mode(part_of_output,axis=1)[0]\n",
    "    #print(max_of_part)\n",
    "    percentage = np.sum(np.equal(max_of_part ,test_output_expected)*1.0)\n",
    "    #print(percentage)\n",
    "    accuracy_over_time[idx] = percentage*100/total_samples\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[10.1  10.1  22.81 44.41 86.6  91.14 96.22 96.81 97.37 97.55 97.65 97.7\n",
      " 97.86 97.85 97.9  97.91 97.95 97.98 97.99 98.03 98.05 98.08 98.09 98.07\n",
      " 98.06 98.06 98.08 98.08 98.1  98.1  98.11 98.13 98.15 98.15 98.15 98.15\n",
      " 98.14 98.14 98.15 98.16 98.17 98.17 98.15 98.15 98.14 98.13 98.13 98.13\n",
      " 98.13 98.13 98.13 98.14 98.14 98.14 98.15 98.17 98.17 98.17 98.17 98.16\n",
      " 98.16 98.16 98.16 98.16 98.17 98.18 98.18 98.2  98.19 98.2  98.2  98.2\n",
      " 98.2  98.2  98.2  98.22 98.22 98.22 98.22 98.23 98.23 98.23 98.24 98.24\n",
      " 98.24 98.24 98.24 98.24 98.24 98.24 98.24 98.24 98.24 98.24 98.23 98.22\n",
      " 98.21 98.21 98.21 98.21 98.21 98.21 98.21 98.21 98.21 98.21 98.21 98.21\n",
      " 98.21 98.21 98.21 98.21 98.21 98.21 98.21 98.21 98.21 98.21 98.21 98.21\n",
      " 98.21 98.21 98.2  98.2  98.18 98.18 98.18 98.18 98.18 98.18 98.17 98.17\n",
      " 98.17 98.17 98.17 98.17 98.17 98.17 98.17 98.16 98.16 98.16 98.16 98.16\n",
      " 98.16 98.16 98.17 98.17 98.17 98.17 98.17 98.17 98.17 98.17 98.17 98.17\n",
      " 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16\n",
      " 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.17 98.16\n",
      " 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16\n",
      " 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.15 98.15 98.15\n",
      " 98.15 98.15 98.15 98.15 98.15 98.15 98.15 98.15 98.15 98.15 98.15 98.15\n",
      " 98.15 98.15 98.15 98.15 98.15 98.15 98.15 98.15 98.15 98.15 98.15 98.15\n",
      " 98.15 98.15 98.15 98.15 98.15 98.15 98.15 98.15 98.15 98.15 98.15 98.15\n",
      " 98.15 98.15 98.15 98.15 98.15 98.15 98.14 98.14 98.14 98.14 98.14 98.14\n",
      " 98.14 98.14 98.14 98.14 98.14 98.14 98.14 98.14 98.14 98.14 98.14 98.14\n",
      " 98.14 98.14 98.14 98.14 98.14 98.14 98.14 98.14 98.14 98.14 98.14 98.14\n",
      " 98.14 98.14 98.14 98.14 98.14 98.14 98.14 98.15 98.15 98.15 98.15 98.15\n",
      " 98.15 98.15 98.15 98.15 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16\n",
      " 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16 98.16\n",
      " 98.16 98.16 98.17 98.17 98.17 98.17 98.17 98.17 98.17 98.17 98.18 98.18\n",
      " 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18\n",
      " 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18\n",
      " 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18\n",
      " 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18\n",
      " 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18\n",
      " 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18 98.18\n",
      " 98.18 98.18 98.19 98.19]\n"
     ]
    }
   ],
   "source": [
    "print(accuracy_over_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#print(output_array[199])\n",
    "print(train_output_array.shape)\n",
    "print(train_output_expected.shape)\n",
    "print(test_output_array.shape)\n",
    "print(test_output_expected.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class compute_analysis():\n",
    "    def __init__(self,inp_ch, out_ch, padding='valid', filter_size=(3,3)):\n",
    "        #super.init(compute_analysis,self).__init__()\n",
    "        self.inp_ch = inp_ch\n",
    "        self.out_ch = out_ch\n",
    "        self.filter_size = filter_size\n",
    "        self.padding = padding\n",
    "        self.total_number_of_conv_ops = 0\n",
    "        self.single_conv_ops =0 \n",
    "        self.shape = None\n",
    "        self.output_shape = np.zeros(2)\n",
    "        self.first_ifmap = 1\n",
    "        self.ifmap_sparsity_list = [] #torch.zeros(1)\n",
    "        self.number_of_zero_ifmaps = 0\n",
    "        self.number_of_nonzero_addtions =0\n",
    "        self.layer_dict ={}\n",
    "        self.sparsity_tensor = None\n",
    "        self.conv_counter = nn.Conv2d(inp_ch,out_ch,kernel_size=filter_size).cuda()\n",
    "        self.conv_weight = torch.zeros((out_ch,inp_ch,*filter_size),dtype= torch.float32)+1.0\n",
    "        self.conv_bias = torch.zeros(out_ch)\n",
    "\n",
    "        self.conv_counter.weight = nn.Parameter(self.conv_weight.cuda())\n",
    "        self.conv_counter.bias = nn.Parameter(self.conv_bias.cuda())\n",
    "        \n",
    "    ## This will measure sparsity over input images    \n",
    "    def crude_conv_analysis(self,input_frame,first_frame):\n",
    "        if first_frame == True:\n",
    "            self.shape = input_frame.shape\n",
    "            shape = self.shape\n",
    "            number_of_inp_pix = shape[0]*shape[1]*shape[2]*shape[3]\n",
    "            if self.padding == 'valid':\n",
    "                self.output_shape[0] = int(shape[2]-self.filter_size[0]+1)\n",
    "                self.output_shape[1] = int(shape[3]-self.filter_size[1]+1)\n",
    "            else:\n",
    "                self.output_shape[0] = int(shape[2])\n",
    "                self.output_shape[1] = int(shape[3])\n",
    "            \n",
    "            self.total_number_of_conv_ops = shape[0]*(self.output_shape[0])*(self.output_shape[1])*self.out_ch*self.filter_size[0]*self.filter_size[1]*shape[1]#self.inp_ch\n",
    "            self.single_conv_ops = (self.output_shape[0])*(self.output_shape[1])*self.out_ch*self.filter_size[0]*self.filter_size[1]*self.inp_ch\n",
    "            number_of_mult_in_conv_per_op = self.filter_size[0]*self.filter_size[1]*self.inp_ch\n",
    "            for i in range(shape[0]):\n",
    "                for j in range(shape[1]):\n",
    "                    for id1 in range(int(self.output_shape[0])):\n",
    "                        for id2 in range(int(self.output_shape[1])):\n",
    "                            if self.first_ifmap == 1:\n",
    "                                number_of_on_fields = torch.sum(input_frame[i,j,id1:id1+self.filter_size[0],id2:id2+self.filter_size[1]])\n",
    "                                self.number_of_nonzero_addtions += number_of_on_fields*self.out_ch\n",
    "                                if number_of_on_fields ==0:\n",
    "                                    self.number_of_zero_ifmaps+=1\n",
    "                                #print('before',self.ifmap_sparsity_list,torch.div(number_of_on_fields,number_of_mult_in_conv_per_op))\n",
    "                                self.ifmap_sparsity_list.append(torch.div(number_of_on_fields,number_of_mult_in_conv_per_op).detach().numpy().item(0))\n",
    "                                #print('after',self.ifmap_sparsity_list)\n",
    "\n",
    "                                #self.ifmap_sparsity_list.append(torch.div(number_of_on_fields,number_of_mult_in_conv_per_op))\n",
    "                                #print(torch.div(number_of_on_fields,number_of_mult_in_conv_per_op))#self.ifmap_sparsity_list)\n",
    "                                self.first_ifmap = 0\n",
    "                            else:\n",
    "                                number_of_on_fields = torch.sum(input_frame[i,j,id1:id1+self.filter_size[0],id2:id2+self.filter_size[1]])\n",
    "                                self.number_of_nonzero_addtions += number_of_on_fields*self.out_ch\n",
    "                                if number_of_on_fields ==0:\n",
    "                                    self.number_of_zero_ifmaps+=1\n",
    "                                #print(torch.div(number_of_on_fields,number_of_mult_in_conv_per_op))#self.ifmap_sparsity_list)\n",
    "                                #print('before',self.ifmap_sparsity_list)\n",
    "\n",
    "                                self.ifmap_sparsity_list.append(torch.div(number_of_on_fields,number_of_mult_in_conv_per_op).detach().numpy().item(0))\n",
    "                                #print('after',self.ifmap_sparsity_list)\n",
    "\n",
    "                                #self.ifmap_sparsity_list = torch.cat((self.ifmap_sparsity_list[0],torch.div(number_of_on_fields,number_of_mult_in_conv_per_op)),0)\n",
    "        \n",
    "        else:\n",
    "            self.shape = input_frame.shape\n",
    "            shape = self.shape\n",
    "            number_of_inp_pix = shape[0]*shape[1]*shape[2]*shape[3]\n",
    "\n",
    "            self.total_number_of_conv_ops += shape[0]*(self.output_shape[0])*(self.output_shape[1])*self.out_ch*self.filter_size[0]*self.filter_size[1]*shape[1]#self.inp_ch\n",
    "            number_of_mult_in_conv_per_op = self.filter_size[0]*self.filter_size[1]*self.inp_ch\n",
    "            for i in range(shape[0]):\n",
    "                for j in range(shape[1]):\n",
    "                    for id1 in range(int(self.output_shape[0])):\n",
    "                        for id2 in range(int(self.output_shape[1])):\n",
    "                            number_of_on_fields = torch.sum(input_frame[i,j,id1:id1+self.filter_size[0],id2:id2+self.filter_size[1]])\n",
    "                            self.number_of_nonzero_addtions += number_of_on_fields*self.out_ch\n",
    "                            if number_of_on_fields ==0:\n",
    "                                self.number_of_zero_ifmaps+=1\n",
    "                            #print(torch.div(number_of_on_fields,number_of_mult_in_conv_per_op))#self.ifmap_sparsity_list)\n",
    "                            #self.ifmap_sparsity_list = torch.cat((self.ifmap_sparsity_list,torch.div(number_of_on_fields,number_of_mult_in_conv_per_op)),0)\n",
    "                            self.ifmap_sparsity_list.append(torch.div(number_of_on_fields,number_of_mult_in_conv_per_op).detach().numpy().item(0))\n",
    "\n",
    "                            #self.ifmap_sparsity_list = torch.cat((self.ifmap_sparsity_list[0],torch.div(number_of_on_fields,number_of_mult_in_conv_per_op)),0)\n",
    "                            \n",
    "        \n",
    "\n",
    "            \n",
    "        return 0\n",
    "    def conv_analysis(self,input_frame,first_frame):\n",
    "        if first_frame == True:\n",
    "            self.shape = input_frame.shape\n",
    "            shape = self.shape\n",
    "            number_of_inp_pix = shape[0]*shape[1]*shape[2]*shape[3]\n",
    "            if self.padding == 'valid':\n",
    "                self.output_shape[0] = int(shape[2]-self.filter_size[0]+1)\n",
    "                self.output_shape[1] = int(shape[3]-self.filter_size[1]+1)\n",
    "            else:\n",
    "                self.output_shape[0] = int(shape[2])\n",
    "                self.output_shape[1] = int(shape[3])\n",
    "\n",
    "            self.total_number_of_conv_ops = shape[0]*(self.output_shape[0])*(self.output_shape[1])*self.out_ch*self.filter_size[0]*self.filter_size[1]*shape[1]#self.inp_ch\n",
    "            self.single_conv_ops = (self.output_shape[0])*(self.output_shape[1])*self.out_ch*self.filter_size[0]*self.filter_size[1]*self.inp_ch\n",
    "            number_of_mult_in_conv_per_op = self.filter_size[0]*self.filter_size[1]*self.inp_ch*self.out_ch\n",
    "            #self.hist_bins = np.arange(number_of_mult_in_conv_per_op+1)\n",
    "            number_of_on_fields = self.conv_counter(input_frame).detach().cpu()\n",
    "            self.sparsity_tensor= torch.histc(number_of_on_fields,bins = number_of_mult_in_conv_per_op+1,max = number_of_mult_in_conv_per_op,min=0)\n",
    "            self.number_of_nonzero_addtions += torch.sum(number_of_on_fields).detach().item()\n",
    "            self.number_of_zero_ifmaps += torch.sum(torch.eq(number_of_on_fields,0)*1.0).detach().item()\n",
    "            #self.sparsity_tensor = number_of_on_fields\n",
    "            #self.sparsity_tensor = torch.div(number_of_on_fields,number_of_mult_in_conv_per_op).detach()\n",
    "        \n",
    "        else:\n",
    "            self.shape = input_frame.shape\n",
    "            shape = self.shape\n",
    "            number_of_inp_pix = shape[0]*shape[1]*shape[2]*shape[3]\n",
    "\n",
    "            self.total_number_of_conv_ops += shape[0]*(self.output_shape[0])*(self.output_shape[1])*self.out_ch*self.filter_size[0]*self.filter_size[1]*shape[1]#self.inp_ch\n",
    "            number_of_mult_in_conv_per_op = self.filter_size[0]*self.filter_size[1]*self.inp_ch*self.out_ch\n",
    "            number_of_on_fields = self.conv_counter(input_frame).detach().cpu()\n",
    "            self.number_of_nonzero_addtions += torch.sum(number_of_on_fields).detach().item()\n",
    "            self.number_of_zero_ifmaps += torch.sum(torch.eq(number_of_on_fields,0)*1.0).detach().item()\n",
    "            self.sparsity_tensor.add_(torch.histc(number_of_on_fields,bins = number_of_mult_in_conv_per_op+1,max = number_of_mult_in_conv_per_op,min=0))\n",
    "\n",
    "            #self.sparsity_tensor = torch.cat((self.sparsity_tensor,number_of_on_fields),0)\n",
    "            #self.sparsity_tensor = torch.cat((self.sparsity_tensor,torch.div(number_of_on_fields,number_of_mult_in_conv_per_op).detach()),0)\n",
    "\n",
    "            \n",
    "        return 0\n",
    "    \n",
    "    \n",
    "    \n",
    "    def summary(self):\n",
    "        self.layer_dict  = {'total_pure_ann_ops':self.total_number_of_conv_ops, 'single_ann_ops':self.single_conv_ops, 'ifmap_sparsity_list':self.ifmap_sparsity_list,'number_of_zero_ifmaps':self.number_of_zero_ifmaps,'total_nonzero_ops':self.number_of_nonzero_addtions,'sparsity_tensor':torch.flatten(self.sparsity_tensor)}\n",
    "        return self.layer_dict\n",
    "        \n",
    "        \n",
    "    def fc_analysis(self,input_frame):\n",
    "        return 0\n",
    "\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_timesteps = 200\n",
    "conv1_analysis = compute_analysis(1, 8)\n",
    "conv2_analysis = compute_analysis(8, 16)\n",
    "conv3_analysis = compute_analysis(16, 20,(5,5))\n",
    "test_output_expected = np.zeros((test_size,1))\n",
    "test_output_array = np.zeros((test_size,number_of_timesteps))\n",
    "for batch_idx, (data, target) in enumerate(snn_test_loader):\n",
    "    start = batch_idx*train_batch_size\n",
    "    end = batch_idx*train_batch_size + data.shape[0]\n",
    "    data = data.cuda()\n",
    "    \n",
    "    spike_pot = torch.zeros_like(data).detach().cuda()\n",
    "    sp1_np = np.zeros((data.shape[0],8,26,26))\n",
    "    sp1_out = torch.tensor(sp1_np, requires_grad=False, dtype=torch.float32).cuda()\n",
    "    sp2_np = np.zeros((data.shape[0],16,24,24))\n",
    "    sp2_out = torch.tensor(sp2_np, requires_grad=False, dtype=torch.float32).cuda()\n",
    "    sp3_np = np.zeros((data.shape[0],20,8,8))\n",
    "    sp3_out = torch.tensor(sp3_np, requires_grad=False, dtype=torch.float32).cuda()\n",
    "    sp_fc1_np = np.zeros((data.shape[0],50))\n",
    "    sp_fc1_out = torch.tensor(sp_fc1_np, requires_grad=False, dtype=torch.float32).cuda()\n",
    "    sp_fc2_np = np.zeros((data.shape[0],10))\n",
    "    sp_fc2_out = torch.tensor(sp_fc2_np, requires_grad=False, dtype=torch.float32).cuda()\n",
    "    \n",
    "    for i in range(number_of_timesteps):\n",
    "        #print(i)\n",
    "        first_frame = True if i ==0 else False\n",
    "        spike_pot.add_(data)\n",
    "\n",
    "        spike_frame = torch.gt(spike_pot,1.0)*1.0 #F.threshold_(spike_pot,1,0).sign() #\n",
    "        spike_pot.sub_(spike_frame)\n",
    "        conv1_analysis.conv_analysis(spike_frame,first_frame)\n",
    "        sp1_out.add_(spikenet_1(spike_frame).detach())\n",
    "        sp2_in = torch.gt(sp1_out,1.0)*1.0 #F.threshold_(sp1_out, 1, 0) #\n",
    "        sp1_out.sub_(sp2_in)\n",
    "        conv2_analysis.conv_analysis(sp2_in,first_frame)\n",
    "\n",
    "        sp2_out.add_(spikenet_2(sp2_in).detach())\n",
    "        sp3_in = torch.gt(sp2_out,1.0)*1.0 #F.threshold_(sp2_out, 1, 0).sign() # \n",
    "        sp2_out.sub_(sp3_in)\n",
    "        conv3_analysis.conv_analysis(sp3_in,first_frame)\n",
    "\n",
    "        sp3_in_red = max_pool2(sp3_in)\n",
    "        sp3_out.add_(spikenet_3(sp3_in_red).detach())\n",
    "        sp4_in = torch.gt(sp3_out,1.0)*1.0 #F.threshold_(sp3_out, 1, 0).sign() #\n",
    "        sp3_out.sub_(sp4_in)\n",
    "        sp4_in_red = max_pool3(sp4_in)\n",
    "        sp4_in_red_flat = torch.flatten(sp4_in_red,start_dim=1)\n",
    "        sp_fc1_out.add_(spikenet_fc_1(sp4_in_red_flat).detach())\n",
    "        sp_fc2_in = torch.gt(sp_fc1_out,1.0)*1.0 #F.threshold_(sp_fc1_out, 1, 0).sign() #\n",
    "        sp_fc1_out.sub_(sp_fc2_in)\n",
    "        sp_fc2_out.add_(spikenet_fc_2(sp_fc2_in).detach())\n",
    "        sof_pot = F.softmax(sp_fc2_out).detach()\n",
    "        max_numbers = torch.argmax(sof_pot,dim=1)\n",
    "        print(torch.cuda.memory_allocated(device=0))\n",
    "        torch.cuda.empty_cache()\n",
    "    if(batch_idx==5):\n",
    "        break\n",
    "    else:\n",
    "        print('batch_id', batch_idx)\n",
    "\n",
    "\n",
    "\n",
    "print(conv1_analysis.summary())  \n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "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.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}