{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PnRi7rdVuM4g"
   },
   "source": [
    "# Prolegomena\n",
    "\n",
    "Depending on whether you are running locally or on google colab uncomment the appropriate PATH cell.\n",
    "By default the local cell is uncommented."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F6QqL3ac1-nj"
   },
   "source": [
    " ## Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "executionInfo": {
     "elapsed": 4105,
     "status": "ok",
     "timestamp": 1632315180386,
     "user": {
      "displayName": "Janko Petkovic",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjT3sBznJCqbsAowUgTCcTHH6omzrINDkNEwAiN=s64",
      "userId": "02272381216738732698"
     },
     "user_tz": -120
    },
    "id": "cNdnu9WaIPKn"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "torch.manual_seed(42)\n",
    "\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torch.nn import CrossEntropyLoss\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "np.random.seed(0)\n",
    "import pandas as pd\n",
    "\n",
    "from numpy.random import default_rng\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "from PIL import Image\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Local run: PATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "PATH_TO_ROOT = \".\"\n",
    "PATH_TO_DATASETS = os.path.join(PATH_TO_ROOT, \"datasets\")\n",
    "PATH_TO_STATE_DICTS = os.path.join(PATH_TO_ROOT, \"trained_models\")\n",
    "\n",
    "from modules.utils import Trainer\n",
    "from modules.models import DFC_LeNet_5, Deep_RetiNet\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "V8C6xMfy2EHy"
   },
   "source": [
    "## Google Colab run: mount drive, PATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 193,
     "status": "ok",
     "timestamp": 1632315245598,
     "user": {
      "displayName": "Janko Petkovic",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjT3sBznJCqbsAowUgTCcTHH6omzrINDkNEwAiN=s64",
      "userId": "02272381216738732698"
     },
     "user_tz": -120
    },
    "id": "tJNgomig9wuQ",
    "outputId": "7e096b4b-aba6-4e5e-b881-b8f7f09a2dfb"
   },
   "outputs": [],
   "source": [
    "# from google.colab import drive\n",
    "# import sys\n",
    "\n",
    "# drive.mount('./mnt')\n",
    "# !ls mnt/MyDrive/Code/ -l\n",
    "\n",
    "# PATH_TO_ROOT = \"mnt/MyDrive/Code/RetiNet\"\n",
    "# PATH_TO_DATASETS = os.path.join(PATH_TO_ROOT, \"datasets\")\n",
    "# PATH_TO_STATE_DICTS = os.path.join(PATH_TO_ROOT, \"trained_models\")\n",
    "\n",
    "# sys.path.append(PATH_TO_ROOT)\n",
    "# from modules.utils import Trainer\n",
    "# from modules.models import DFC_LeNet_5, Deep_RetiNet\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 19,
     "status": "ok",
     "timestamp": 1632315198933,
     "user": {
      "displayName": "Janko Petkovic",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjT3sBznJCqbsAowUgTCcTHH6omzrINDkNEwAiN=s64",
      "userId": "02272381216738732698"
     },
     "user_tz": -120
    },
    "id": "zzY0KC0e9yW5",
    "outputId": "4ecc13d8-fc29-47c5-ff29-7c12a9f55f73"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using: cpu\n"
     ]
    }
   ],
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(f'Using: {device}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testset transforming methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "executionInfo": {
     "elapsed": 13,
     "status": "ok",
     "timestamp": 1632315198934,
     "user": {
      "displayName": "Janko Petkovic",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjT3sBznJCqbsAowUgTCcTHH6omzrINDkNEwAiN=s64",
      "userId": "02272381216738732698"
     },
     "user_tz": -120
    },
    "id": "sgZfb01ZufdR"
   },
   "outputs": [],
   "source": [
    "class ShiftMean(object):\n",
    "  def __init__(self, mu):\n",
    "    self.mu = mu\n",
    "  \n",
    "  def collate(self, batch):\n",
    "    imgs = []\n",
    "    labels = []\n",
    "\n",
    "    for pair in batch:\n",
    "      shifted = pair[0] - self.mu\n",
    "      imgs.append(shifted)\n",
    "      labels.append(pair[1])\n",
    "\n",
    "    imgs = torch.stack(imgs, dim=0)\n",
    "    labels = torch.tensor(labels)\n",
    "\n",
    "    return [imgs, labels]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "executionInfo": {
     "elapsed": 12,
     "status": "ok",
     "timestamp": 1632315198935,
     "user": {
      "displayName": "Janko Petkovic",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjT3sBznJCqbsAowUgTCcTHH6omzrINDkNEwAiN=s64",
      "userId": "02272381216738732698"
     },
     "user_tz": -120
    },
    "id": "2Fnja54su18R"
   },
   "outputs": [],
   "source": [
    "class ScaleVar(object):\n",
    "  def __init__(self, sigma):\n",
    "    self.sigma = sigma\n",
    "  \n",
    "  def collate(self, batch):\n",
    "    imgs = []\n",
    "    labels = []\n",
    "\n",
    "    for pair in batch:\n",
    "      mean = pair[0].mean()\n",
    "      scaled = (pair[0]-mean)/self.sigma + mean\n",
    "\n",
    "      imgs.append(scaled)\n",
    "      labels.append(pair[1])\n",
    "\n",
    "    imgs = torch.stack(imgs, dim=0)\n",
    "    labels = torch.tensor(labels)\n",
    "\n",
    "    return [imgs, labels]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "J4XNYGrY63S0"
   },
   "source": [
    "# Sweeps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Input panel "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "executionInfo": {
     "elapsed": 316,
     "status": "ok",
     "timestamp": 1632322342511,
     "user": {
      "displayName": "Janko Petkovic",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjT3sBznJCqbsAowUgTCcTHH6omzrINDkNEwAiN=s64",
      "userId": "02272381216738732698"
     },
     "user_tz": -120
    },
    "id": "FKBVq-u_SbHk"
   },
   "outputs": [],
   "source": [
    "\n",
    "dataset_name = \"MNIST\" #@param          MNIST - FashionMNIST - SVHN \n",
    "\n",
    "retinic_kernel_size =  7#@param          retinic kernel size\n",
    "depth = 3 #@param\n",
    "\n",
    "rks = retinic_kernel_size\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Testset loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "executionInfo": {
     "elapsed": 10,
     "status": "ok",
     "timestamp": 1632322342512,
     "user": {
      "displayName": "Janko Petkovic",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjT3sBznJCqbsAowUgTCcTHH6omzrINDkNEwAiN=s64",
      "userId": "02272381216738732698"
     },
     "user_tz": -120
    },
    "id": "sWnGT5FA2uiZ"
   },
   "outputs": [],
   "source": [
    "IL = { \n",
    "    \"MNIST\" : datasets.MNIST,\n",
    "    \"FashionMNIST\" : datasets.FashionMNIST,\n",
    "    \"SVHN\" : datasets.SVHN\n",
    "}\n",
    "\n",
    "\n",
    "if dataset_name == \"SVHN\":\n",
    "    # pytorch does not use the same standards for all\n",
    "    # datasets for some reasons I dont know why\n",
    "    path = os.path.join(PATH_TO_DATASETS, \"SVHN\")\n",
    "\n",
    "    testset = IL[dataset_name](\n",
    "        path,\n",
    "        download = True,\n",
    "        split = \"test\",\n",
    "        transform = transforms.ToTensor()\n",
    "    )\n",
    "\n",
    "else:\n",
    "    transform = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Pad(2)\n",
    "    ])\n",
    "\n",
    "    testset = IL[dataset_name](\n",
    "        PATH_TO_DATASETS,\n",
    "        download = True,\n",
    "        train = False,\n",
    "        transform = transform\n",
    "    )\n",
    " \n",
    "in_channels = testset[0][0].shape[0]\n",
    "test_size = testset.__len__()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Model istantiation and weight loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 312,
     "status": "ok",
     "timestamp": 1632322342815,
     "user": {
      "displayName": "Janko Petkovic",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjT3sBznJCqbsAowUgTCcTHH6omzrINDkNEwAiN=s64",
      "userId": "02272381216738732698"
     },
     "user_tz": -120
    },
    "id": "cAi54tRA4jO_",
    "outputId": "a45a30f9-b532-4a16-e500-533d9eebe620"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retinet_path = os.path.join(\n",
    "    PATH_TO_STATE_DICTS, dataset_name, f\"trained_Deep_RetiNet_d{depth}_rks{rks}_state_dict.pt\") \n",
    "\n",
    "lenet_path = os.path.join(\n",
    "    PATH_TO_STATE_DICTS, dataset_name, f\"trained_DFC_LeNet_5_state_dict.pt\")\n",
    "\n",
    "\n",
    "retinet = Deep_RetiNet(depth, rks, in_channels).to(device).eval()\n",
    "retinet.load_state_dict(torch.load(retinet_path, map_location = device))\n",
    "\n",
    "lenet = DFC_LeNet_5(in_channels).to(device).eval()\n",
    "lenet.load_state_dict(torch.load(lenet_path, map_location = device))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sweep loops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 116318,
     "status": "ok",
     "timestamp": 1632322459426,
     "user": {
      "displayName": "Janko Petkovic",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjT3sBznJCqbsAowUgTCcTHH6omzrINDkNEwAiN=s64",
      "userId": "02272381216738732698"
     },
     "user_tz": -120
    },
    "id": "XG1ArGmCwyef",
    "outputId": "7796ce0b-e562-46a2-d5e5-7e32d0230a07"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 29%|██▊       | 6/21 [00:15<00:37,  2.52s/it]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-11-0b8f6d094409>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     15\u001b[0m     \u001b[0mbatch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlabels\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlabels\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 17\u001b[0;31m     \u001b[0mout_lenet\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlenet\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     18\u001b[0m     \u001b[0mout_retinet\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mretinet\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     19\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    887\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    888\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 889\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    890\u001b[0m         for hook in itertools.chain(\n\u001b[1;32m    891\u001b[0m                 \u001b[0m_global_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Code/retinet/modules/models/dfclenet5.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     41\u001b[0m     \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mc1\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     42\u001b[0m     \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdrop\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 43\u001b[0;31m     \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtanh\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     44\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     45\u001b[0m     \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0ms2\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    887\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    888\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 889\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    890\u001b[0m         for hook in itertools.chain(\n\u001b[1;32m    891\u001b[0m                 \u001b[0m_global_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.9/site-packages/torch/nn/modules/activation.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    357\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    358\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 359\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtanh\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    360\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    361\u001b[0m \u001b[0;32mclass\u001b[0m \u001b[0mSiLU\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mModule\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "## MEAN SWEEP LOOP:   mean = i/10\n",
    "\n",
    "results_mean = []\n",
    "\n",
    "for i in tqdm(range(-20,21,2)):\n",
    "  mean = i/10\n",
    "\n",
    "  transform = ShiftMean(mean)\n",
    "\n",
    "  testloader = DataLoader(\n",
    "      testset, batch_size = test_size,\n",
    "      collate_fn = transform.collate)\n",
    "\n",
    "  for batch, labels in testloader:\n",
    "    batch, labels = batch.to(device), labels.to(device)\n",
    "\n",
    "    out_lenet = lenet(batch)\n",
    "    out_retinet = retinet(batch)\n",
    "\n",
    "    acc_lenet = (labels == out_lenet.argmax(axis=1)).sum()/test_size\n",
    "    acc_retinet = (labels == out_retinet.argmax(axis=1)).sum()/test_size\n",
    "  \n",
    "  results_mean.append([mean, acc_lenet.item(), acc_retinet.item()])\n",
    "\n",
    "\n",
    "# DEVIATION SWEEP LOOP:    deviation = i/10\n",
    "\n",
    "results_var = []\n",
    "\n",
    "for i in tqdm(range(1,41,2)):\n",
    "  var = i/10\n",
    "\n",
    "  transform = ScaleVar(var)\n",
    "\n",
    "  testloader = DataLoader(\n",
    "      testset, batch_size = test_size,\n",
    "      collate_fn = transform.collate)\n",
    "\n",
    "  for batch, labels in testloader:\n",
    "    batch, labels = batch.to(device), labels.to(device)\n",
    "\n",
    "    out_lenet = lenet(batch)\n",
    "    out_retinet = retinet(batch)\n",
    "\n",
    "    acc_lenet = (labels == out_lenet.argmax(axis=1)).sum()/test_size\n",
    "    acc_retinet = (labels == out_retinet.argmax(axis=1)).sum()/test_size\n",
    "  \n",
    "  results_var.append([var, acc_lenet.item(), acc_retinet.item()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results previews"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 398
    },
    "executionInfo": {
     "elapsed": 932,
     "status": "ok",
     "timestamp": 1632322460343,
     "user": {
      "displayName": "Janko Petkovic",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjT3sBznJCqbsAowUgTCcTHH6omzrINDkNEwAiN=s64",
      "userId": "02272381216738732698"
     },
     "user_tz": -120
    },
    "id": "6LcHv--Jxuwh",
    "outputId": "7a83fd05-4bd1-45e9-ad69-ccdb9dbd2cc8"
   },
   "outputs": [],
   "source": [
    "printable_mean = np.array(results_mean)\n",
    "printable_var = np.array(results_var)\n",
    "\n",
    "fig = plt.figure(figsize=(18,5))\n",
    "tight = fig.tight_layout(w_pad=40)\n",
    "fig.set_tight_layout(tight)\n",
    "\n",
    "fig1 = fig.add_subplot(1,2,1)\n",
    "fig1.set_title(\"Accuracy versus $\\mu$ sweep\", size=20, pad=20)\n",
    "\n",
    "fig1.set_xlabel(\"$\\mu$\", fontsize=15, labelpad=15)\n",
    "fig1.set_ylim([0,1])\n",
    "fig1.set_ylabel(\"Accuracy\", fontsize=15, labelpad=15)\n",
    "\n",
    "fig1.tick_params(axis='x', labelsize=15, pad=10)\n",
    "fig1.tick_params(axis='y', labelsize=15, pad=10)\n",
    "\n",
    "plt.scatter(printable_mean[:,0], printable_mean[:,1], \n",
    "            marker='o', label=f'LeNet_5')\n",
    "\n",
    "plt.scatter(printable_mean[:,0], printable_mean[:,2],\n",
    "            marker='o', label=f'Deep_RetiNet_d{depth}_rks{rks}')\n",
    "plt.legend(loc=3, fontsize=15)\n",
    "\n",
    "fig2 = fig.add_subplot(1,2,2)\n",
    "fig2.set_title(\"Accuracy versus $\\sigma$ sweep\", fontsize = 20, pad=20)\n",
    "fig2.set_xlabel(\"$\\sigma$\", fontsize = 15, labelpad=15)\n",
    "fig2.set_ylim([0,1])\n",
    "fig2.set_ylabel(\"Accuracy\", fontsize = 15, labelpad=15)\n",
    "\n",
    "\n",
    "fig2.tick_params(axis='x', labelsize=15, pad=10)\n",
    "fig2.tick_params(axis='y', labelsize=15, pad=10)\n",
    "\n",
    "plt.scatter(printable_var[:,0], printable_var[:,1], \n",
    "            marker='o', label='LeNet_5')\n",
    "plt.scatter(printable_var[:,0], printable_var[:,2], \n",
    "            marker='o', label=f'Deep_RetiNet_d{depth}_rks{rks}')\n",
    "plt.legend(loc=3, fontsize=15)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Output .csv files generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 14,
     "status": "ok",
     "timestamp": 1632322460344,
     "user": {
      "displayName": "Janko Petkovic",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjT3sBznJCqbsAowUgTCcTHH6omzrINDkNEwAiN=s64",
      "userId": "02272381216738732698"
     },
     "user_tz": -120
    },
    "id": "WT3iXj_Vxvp_"
   },
   "outputs": [],
   "source": [
    "df_results_mean = pd.DataFrame(elem for elem in results_mean)\n",
    "df_results_var = pd.DataFrame(elem for elem in results_var)\n",
    "\n",
    "\n",
    "df_results_mean.rename(\n",
    "    columns = {0 : \"mean_off\", 1 : \"acc_lenet\", 2 :\"acc_retinet\"},\n",
    "    inplace = True)\n",
    "\n",
    "df_results_var.rename(\n",
    "    columns = {0 : \"var_resc\", 1 : \"acc_lenet\", 2 :\"acc_retinet\"},\n",
    "    inplace = True)\n",
    "\n",
    "\n",
    "\n",
    "mean_file_name = f\"accuracy_sweep_Deep_RetiNet_d{depth}_rks{rks}_vs_mean.csv\"\n",
    "PATH_TO_SAVE_FILE = os.path.join(\n",
    "    PATH_TO_ROOT, f\"results/accuracy_sweeps/{dataset_name}\", mean_file_name)\n",
    "\n",
    "df_results_mean.to_csv(PATH_TO_SAVE_FILE, index=False)\n",
    "\n",
    "\n",
    "var_file_name = f\"accuracy_sweep_Deep_RetiNet_d{depth}_rks{rks}_vs_var.csv\"\n",
    "PATH_TO_SAVE_FILE = os.path.join(\n",
    "    PATH_TO_ROOT, f\"results/accuracy_sweeps/{dataset_name}\", var_file_name)\n",
    "\n",
    "df_results_var.to_csv(PATH_TO_SAVE_FILE, index=False)\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "accuracy_sweep",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
