notebooks/domino_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "IPython.notebook.set_autosave_interval(10000)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Autosaving every 10 seconds\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2  \n",
    "%autosave 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch, torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import torch.backends.cudnn as cudnn\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from PIL import Image\n",
    "import random, os, sys, argparse\n",
    "from pathlib import Path\n",
    "from tqdm import tqdm\n",
    "import pickle\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.path.insert(0,'/jet/home/nmurali/asc170022p/nmurali/projects/misc/TMLR23_Dynamics_of_Spurious_Features/scripts')\n",
    "from models import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# user hyperparams\n",
    "model = 'resnet18' # resnet18, vgg16, densenet121\n",
    "train_csv = '/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/experiments/toy_expts/domino_expts/top_svhn_bot_fmnist_2class_ro_1p0.csv'\n",
    "test_csv = '/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/experiments/toy_expts/domino_expts/top_blank_bot_fmnist_2class_ro_1p0.csv'\n",
    "expt_name = f'temp_{int(random.random()*100000)}'# f'resnet18_top_mnist_bot_kmnist_{int(random.random()*100000)}'\n",
    "save_dir = '/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/experiments/toy_expts/domino_expts/output/'\n",
    "\n",
    "num_epochs = 50\n",
    "lr = 0.1\n",
    "seed = 10\n",
    "num_ch = 3 # num of channels in image\n",
    "num_embs = 2000 # 1500 for mnist, 10k for cifar10\n",
    "K = 29 # K neighbours\n",
    "num_test_imgs = 20 # num of test images for plotting PD\n",
    "lp_norm = 1 # for computing KNN\n",
    "knn_pos_thresh = 0.5\n",
    "knn_neg_thresh = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting the seed\n",
    "np.random.seed(seed)\n",
    "random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "torch.cuda.manual_seed_all(seed)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train = pd.read_csv(train_csv)\n",
    "df_test = pd.read_csv(test_csv).sample(500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ToyDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, class_field, transform, csv_path=None, df=None):\n",
    "        self.class_field = class_field        \n",
    "        self.transform = transform\n",
    "        if df is not None:\n",
    "            self.df = df\n",
    "        else:\n",
    "            self.csvpath = csvpath\n",
    "            self.df = pd.read_csv(self.csvpath)\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img_path = str(self.df.iloc[idx][\"path\"])\n",
    "        lab = int(self.df.iloc[idx][self.class_field])\n",
    "        img = Image.open(img_path)\n",
    "        if self.transform is not None:\n",
    "            img = self.transform(img)  \n",
    "\n",
    "        return {\"img\":img, \"lab\":lab, \"idx\":idx, \"file_name\" : img_path}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "trans = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.ToTensor(),\n",
    "    torchvision.transforms.Resize((64,32)),\n",
    "])\n",
    "trainset = ToyDataset(class_field=['bottom'], transform=trans, df=df_train)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=256, shuffle=True, num_workers=2)\n",
    "testset = ToyDataset(class_field=['bottom'], transform=trans, df=df_test)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=256, shuffle=True, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==> Building model..\n",
      "\n",
      "Epoch: 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [03:50<00:00,  6.08s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 2.503 | Acc: 81.792% (7852/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:11<00:00,  5.59s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 6.535 | Acc: 59.200% (296/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.74it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.130 | Acc: 96.031% (9219/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.146 | Acc: 95.200% (476/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.82it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.095 | Acc: 96.896% (9302/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.20it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.093 | Acc: 96.200% (481/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.076 | Acc: 97.427% (9353/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.15it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.072 | Acc: 97.800% (489/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 4\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.84it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.067 | Acc: 97.688% (9378/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.065 | Acc: 98.200% (491/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.94it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.057 | Acc: 98.031% (9411/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.04it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.059 | Acc: 98.000% (490/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 6\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.92it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.054 | Acc: 98.312% (9438/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.29it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.064 | Acc: 98.000% (490/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.054 | Acc: 98.156% (9423/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.30it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.066 | Acc: 97.600% (488/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 8\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.93it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.045 | Acc: 98.448% (9451/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.21it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.056 | Acc: 98.400% (492/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 9\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.95it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.042 | Acc: 98.573% (9463/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.27it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.053 | Acc: 98.400% (492/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.79it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.036 | Acc: 98.812% (9486/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.048 | Acc: 98.400% (492/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 11\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.85it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.033 | Acc: 98.917% (9496/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.32it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.061 | Acc: 98.000% (490/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 12\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.030 | Acc: 99.021% (9506/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.34it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.054 | Acc: 98.400% (492/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 13\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.82it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.030 | Acc: 98.979% (9502/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.29it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.036 | Acc: 98.800% (494/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 14\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:07<00:00,  5.36it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.027 | Acc: 99.083% (9512/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  2.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.040 | Acc: 98.000% (490/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 15\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.83it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.023 | Acc: 99.198% (9523/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.42it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.046 | Acc: 98.200% (491/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 16\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.73it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.023 | Acc: 99.146% (9518/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.050 | Acc: 98.600% (493/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 17\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.90it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.021 | Acc: 99.323% (9535/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.38it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.034 | Acc: 98.800% (494/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 18\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.90it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.019 | Acc: 99.292% (9532/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.22it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.038 | Acc: 98.600% (493/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 19\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.81it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.017 | Acc: 99.448% (9547/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.28it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.034 | Acc: 98.400% (492/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 20\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  6.00it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.015 | Acc: 99.490% (9551/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.31it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.024 | Acc: 99.400% (497/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 21\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.93it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.014 | Acc: 99.510% (9553/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.25it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.016 | Acc: 99.200% (496/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 22\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.93it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.015 | Acc: 99.469% (9549/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.27it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.021 | Acc: 99.400% (497/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 23\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.93it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.014 | Acc: 99.521% (9554/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.33it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.031 | Acc: 98.600% (493/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 24\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.94it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.013 | Acc: 99.500% (9552/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.33it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.018 | Acc: 99.200% (496/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 25\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.83it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.011 | Acc: 99.667% (9568/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.45it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.015 | Acc: 99.400% (497/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 26\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.94it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.010 | Acc: 99.688% (9570/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.21it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.029 | Acc: 98.800% (494/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 27\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.74it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.010 | Acc: 99.646% (9566/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.35it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.023 | Acc: 99.200% (496/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 28\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.97it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.012 | Acc: 99.604% (9562/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.16it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.016 | Acc: 99.600% (498/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 29\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  6.04it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.007 | Acc: 99.802% (9581/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.29it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.011 | Acc: 99.800% (499/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 30\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.83it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.008 | Acc: 99.729% (9574/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.34it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.012 | Acc: 99.600% (498/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 31\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.96it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.008 | Acc: 99.719% (9573/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.16it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.011 | Acc: 99.800% (499/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 32\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  6.02it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.006 | Acc: 99.844% (9585/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.18it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.011 | Acc: 99.600% (498/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 33\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.95it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.005 | Acc: 99.812% (9582/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.32it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.016 | Acc: 99.600% (498/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 34\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.95it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.007 | Acc: 99.792% (9580/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.10it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.013 | Acc: 99.800% (499/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 35\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.87it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.005 | Acc: 99.854% (9586/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  2.97it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.008 | Acc: 99.800% (499/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 36\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.56it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.004 | Acc: 99.906% (9591/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.25it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.014 | Acc: 99.800% (499/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 37\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.004 | Acc: 99.917% (9592/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.008 | Acc: 99.800% (499/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 38\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.89it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.003 | Acc: 99.958% (9596/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.22it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.012 | Acc: 99.800% (499/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 39\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  6.02it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.005 | Acc: 99.844% (9585/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.22it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.014 | Acc: 99.600% (498/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 40\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.86it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.007 | Acc: 99.760% (9577/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.31it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.009 | Acc: 99.800% (499/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 41\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.99it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.004 | Acc: 99.865% (9587/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.27it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.007 | Acc: 100.000% (500/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 42\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.005 | Acc: 99.781% (9579/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.32it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.009 | Acc: 99.800% (499/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 43\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.004 | Acc: 99.875% (9588/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.25it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.016 | Acc: 99.200% (496/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 44\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.007 | Acc: 99.792% (9580/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.20it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.013 | Acc: 99.800% (499/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 45\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.94it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.006 | Acc: 99.812% (9582/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.29it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.014 | Acc: 99.200% (496/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 46\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.89it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.006 | Acc: 99.812% (9582/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.07it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.014 | Acc: 99.600% (498/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 47\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.90it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.004 | Acc: 99.854% (9586/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.16it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.011 | Acc: 99.400% (497/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 48\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.96it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.005 | Acc: 99.875% (9588/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.07it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.010 | Acc: 99.600% (498/500)\n",
      "Saving..\n",
      "\n",
      "Epoch: 49\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38/38 [00:06<00:00,  5.91it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37 38 Loss: 0.004 | Acc: 99.854% (9586/9600)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00,  3.30it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 2 Loss: 0.009 | Acc: 99.800% (499/500)\n",
      "Saving..\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "best_acc = 0  # best test accuracy\n",
    "\n",
    "\n",
    "# Model\n",
    "print('==> Building model..')\n",
    "if model=='resnet18':\n",
    "    net = ResNet18(num_channels=num_ch)\n",
    "elif model=='vgg16':\n",
    "    net = VGG('VGG16',num_channels=num_ch)\n",
    "elif model=='densenet121':\n",
    "    net = DenseNet121()\n",
    "net.linear = nn.Linear(in_features=1024,out_features=10,bias=True)\n",
    "net = net.to(device)\n",
    "    \n",
    "if device == 'cuda':\n",
    "    net = torch.nn.DataParallel(net)\n",
    "    cudnn.benchmark = True\n",
    "\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(net.parameters(), lr=lr,\n",
    "                      momentum=0.9, weight_decay=5e-4)\n",
    "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)\n",
    "\n",
    "\n",
    "# Training\n",
    "def train(epoch):\n",
    "    print('\\nEpoch: %d' % epoch)\n",
    "    net.train()\n",
    "    train_loss = 0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    for batch_idx, batch in enumerate(tqdm(trainloader)):\n",
    "        inputs = batch['img']\n",
    "        targets = batch['lab']\n",
    "        inputs, targets = inputs.to(device), targets.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        outputs = net(inputs)\n",
    "        loss = criterion(outputs, targets)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        train_loss += loss.item()\n",
    "        _, predicted = outputs.max(1)\n",
    "        total += targets.size(0)\n",
    "        correct += predicted.eq(targets).sum().item()\n",
    "\n",
    "    print(batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'\n",
    "                 % (train_loss/(batch_idx+1), 100.*correct/total, correct, total))\n",
    "\n",
    "\n",
    "def test(epoch):\n",
    "    global best_acc\n",
    "    net.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, batch in enumerate(tqdm(testloader)):\n",
    "            inputs = batch['img']\n",
    "            targets = batch['lab']\n",
    "            inputs, targets = inputs.to(device), targets.to(device)\n",
    "            outputs = net(inputs)\n",
    "            loss = criterion(outputs, targets)\n",
    "\n",
    "            test_loss += loss.item()\n",
    "            _, predicted = outputs.max(1)\n",
    "            total += targets.size(0)\n",
    "            correct += predicted.eq(targets).sum().item()\n",
    "\n",
    "    print(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'\n",
    "                 % (test_loss/(batch_idx+1), 100.*correct/total, correct, total))\n",
    "\n",
    "    # Save checkpoint.\n",
    "    acc = 100.*correct/total\n",
    "#     if acc > best_acc:\n",
    "#         print('Saving..')\n",
    "#         state = {\n",
    "#             'net': net.state_dict(),\n",
    "#             'acc': acc,\n",
    "#             'epoch': epoch,\n",
    "#         }\n",
    "# #         if not os.path.isdir(f'{args['expt_name']}_checkpoint'):\n",
    "# #             os.mkdir('checkpoint')\n",
    "#         torch.save(state, os.path.join(save_dir,f'{expt_name}_ep{epoch}.pt'))\n",
    "#         best_acc = acc\n",
    "    print('Saving..')\n",
    "    state = {\n",
    "        'net': net.state_dict(),\n",
    "        'acc': acc,\n",
    "        'epoch': epoch,\n",
    "    }\n",
    "    torch.save(state, os.path.join(save_dir,f'{expt_name}.pt'))\n",
    "\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    train(epoch)\n",
    "    test(epoch)\n",
    "    scheduler.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Obtain Train Subset Embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DataParallel(\n",
       "  (module): ResNet(\n",
       "    (conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (layer1): Sequential(\n",
       "      (0): BasicBlock(\n",
       "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (shortcut): Sequential()\n",
       "      )\n",
       "      (1): BasicBlock(\n",
       "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (shortcut): Sequential()\n",
       "      )\n",
       "    )\n",
       "    (layer2): Sequential(\n",
       "      (0): BasicBlock(\n",
       "        (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (shortcut): Sequential(\n",
       "          (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (1): BasicBlock(\n",
       "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (shortcut): Sequential()\n",
       "      )\n",
       "    )\n",
       "    (layer3): Sequential(\n",
       "      (0): BasicBlock(\n",
       "        (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (shortcut): Sequential(\n",
       "          (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (1): BasicBlock(\n",
       "        (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (shortcut): Sequential()\n",
       "      )\n",
       "    )\n",
       "    (layer4): Sequential(\n",
       "      (0): BasicBlock(\n",
       "        (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (shortcut): Sequential(\n",
       "          (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (1): BasicBlock(\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (shortcut): Sequential()\n",
       "      )\n",
       "    )\n",
       "    (linear): Linear(in_features=1024, out_features=10, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load checkpoint\n",
    "# expt_name = 'resnet18_top_cifar10_bot_fmnist_64x32_60744_ep6'\n",
    "if model=='resnet18':\n",
    "    net = ResNet18(num_channels=num_ch)\n",
    "    net.linear = nn.Linear(in_features=1024,out_features=10,bias=True)\n",
    "    net = nn.DataParallel(net)\n",
    "elif model=='vgg16':\n",
    "    net = nn.DataParallel(VGG('VGG16',num_channels=num_ch))\n",
    "net.load_state_dict(torch.load(os.path.join(save_dir,f'{expt_name}.pt'))['net'])\n",
    "net.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_resnet18_hooks(net, hook):\n",
    "    net.module.bn1.register_forward_hook(hook)\n",
    "    \n",
    "    net.module.layer1[0].conv1.register_forward_hook(hook)\n",
    "    net.module.layer1[0].conv2.register_forward_hook(hook)\n",
    "    \n",
    "    net.module.layer1[1].conv1.register_forward_hook(hook)\n",
    "    net.module.layer1[1].conv2.register_forward_hook(hook)\n",
    "\n",
    "    net.module.layer2[0].conv1.register_forward_hook(hook)\n",
    "    net.module.layer2[0].conv2.register_forward_hook(hook)\n",
    "    \n",
    "    net.module.layer2[1].conv1.register_forward_hook(hook)\n",
    "    net.module.layer2[1].conv2.register_forward_hook(hook)\n",
    "    \n",
    "    net.module.layer3[0].conv1.register_forward_hook(hook)\n",
    "    net.module.layer3[0].conv2.register_forward_hook(hook)\n",
    "    \n",
    "    net.module.layer3[1].conv1.register_forward_hook(hook)\n",
    "    net.module.layer3[1].conv2.register_forward_hook(hook)\n",
    "    \n",
    "    net.module.layer4[0].conv1.register_forward_hook(hook)\n",
    "    net.module.layer4[0].conv2.register_forward_hook(hook)\n",
    "    \n",
    "    net.module.layer4[1].conv1.register_forward_hook(hook)\n",
    "    net.module.layer4[1].conv2.register_forward_hook(hook)\n",
    "        \n",
    "    return net\n",
    "\n",
    "def add_vgg16_hooks(net, hook):\n",
    "    net.module.features[0].register_forward_hook(hook)\n",
    "    net.module.features[3].register_forward_hook(hook)\n",
    "    net.module.features[7].register_forward_hook(hook)\n",
    "    net.module.features[10].register_forward_hook(hook)\n",
    "    net.module.features[14].register_forward_hook(hook)\n",
    "    net.module.features[17].register_forward_hook(hook)\n",
    "    net.module.features[20].register_forward_hook(hook)\n",
    "    net.module.features[24].register_forward_hook(hook)\n",
    "    net.module.features[27].register_forward_hook(hook)\n",
    "    net.module.features[30].register_forward_hook(hook)\n",
    "    net.module.features[34].register_forward_hook(hook)\n",
    "    net.module.features[37].register_forward_hook(hook)\n",
    "    net.module.features[40].register_forward_hook(hook)        \n",
    "    return net\n",
    "\n",
    "def add_densenet121_hooks(net, hook):\n",
    "    \n",
    "    for idx,layer in enumerate(net.module.dense1):\n",
    "        if idx%2==0:\n",
    "            layer.register_forward_hook(hook)\n",
    "        \n",
    "    for idx,layer in enumerate(net.module.dense2):\n",
    "        if idx%2==0:\n",
    "            layer.register_forward_hook(hook)\n",
    "        \n",
    "    for idx,layer in enumerate(net.module.dense3):\n",
    "        if idx%2==0:\n",
    "            layer.register_forward_hook(hook)\n",
    "        \n",
    "    for idx,layer in enumerate(net.module.dense4):\n",
    "        if idx%2==0:\n",
    "            layer.register_forward_hook(hook)\n",
    "        \n",
    "    return net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_maps = []  # This will be a list of Tensors, each representing a feature map\n",
    "def hook_feat_map(mod, inp, out):\n",
    "#     out = torch.nn.functional.interpolate(out,(24,12))\n",
    "    feature_maps.append(torch.reshape(out, (out.shape[0],-1)))\n",
    "\n",
    "if model=='resnet18':\n",
    "    net = add_resnet18_hooks(net, hook_feat_map)\n",
    "elif model=='vgg16':\n",
    "    net = add_vgg16_hooks(net, hook_feat_map)\n",
    "elif model=='densenet121':\n",
    "    net = add_densenet121_hooks(net, hook_feat_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_cpu(arr):\n",
    "    for idx,x in enumerate(arr):\n",
    "        arr[idx] = x.to('cpu')\n",
    "    return arr\n",
    "\n",
    "def print_memory_profile(s):\n",
    "    # print GPU memory\n",
    "    t = torch.cuda.get_device_properties(0).total_memory\n",
    "    r = torch.cuda.memory_reserved(0)\n",
    "    a = torch.cuda.memory_allocated(0)\n",
    "    print(s)\n",
    "    print(t/1024**3,r/1024**3,a/1024**3)\n",
    "    print('\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2000/2000 [00:08<00:00, 225.82it/s]\n"
     ]
    }
   ],
   "source": [
    "perm = torch.randperm(len(trainset))\n",
    "inds = perm[:num_embs]\n",
    "\n",
    "labs = []\n",
    "samples = torch.empty((0,3,64,32))\n",
    "for i in tqdm(inds):\n",
    "    i = int(i)\n",
    "    labs.append(trainset[i]['lab'])\n",
    "    samples = torch.cat((samples,trainset[i]['img'].unsqueeze(0)))\n",
    "labs = torch.tensor(labs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_resized = samples\n",
    "train_subset = torch.utils.data.TensorDataset(samples_resized,labs)\n",
    "trainloader2 = torch.utils.data.DataLoader(train_subset, batch_size=128, shuffle=True, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  0%|          | 0/1 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial\n",
      "31.7391357421875 4.234375 0.1698150634765625\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:15<00:00, 15.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "After processing Batch\n",
      "31.7391357421875 20.103515625 8.5164475440979\n",
      "\n",
      "\n",
      "After freeing GPU memory\n",
      "31.7391357421875 4.23046875 0.2156662940979004\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# code for saving pkl file of layer embeddings\n",
    "save_path = os.path.join(save_dir,f'{expt_name}.pkl')\n",
    "trainloader2 = torch.utils.data.DataLoader(train_subset, batch_size=30000, shuffle=True, num_workers=2)\n",
    "\n",
    "handle = open(save_path, \"wb\")\n",
    "\n",
    "with torch.no_grad():\n",
    "    net.eval()\n",
    "    for b_idx,batch in enumerate(tqdm(trainloader2)):        \n",
    "        # print GPU memory\n",
    "        print_memory_profile('Initial')\n",
    "        \n",
    "        imgs = batch[0].to('cuda')\n",
    "        labels = batch[1]\n",
    "        \n",
    "        feature_maps = []\n",
    "        out = net(imgs.float())\n",
    "        \n",
    "        info_dict = {'batch_idx':b_idx,'num_batches':len(trainloader2),'feats':feature_maps,'labels':labels}\n",
    "        pickle.dump(info_dict, handle)  \n",
    "        \n",
    "        # print GPU memory\n",
    "        print_memory_profile('After processing Batch')\n",
    "        \n",
    "        # free up GPU memory\n",
    "        del feature_maps, info_dict\n",
    "        torch.cuda.empty_cache()     \n",
    "        \n",
    "        # print GPU memory\n",
    "        print_memory_profile('After freeing GPU memory')\n",
    "        \n",
    "handle.close()\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute PD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute_pd.py\n",
    "\n",
    "def compute_pred_depth(arr):\n",
    "    last = arr[-1]\n",
    "\n",
    "    p_depth = 1\n",
    "    for i in range(len(arr)-1):\n",
    "        ele = arr[-1-(i+1)]\n",
    "        if ele!=last:\n",
    "            p_depth = (len(arr)-(i+1)) + 1\n",
    "            break\n",
    "    \n",
    "    return p_depth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:00<00:00, 356.81it/s]\n"
     ]
    }
   ],
   "source": [
    "perm = torch.randperm(len(testset))\n",
    "inds = perm[:num_test_imgs]\n",
    "\n",
    "labs = []\n",
    "samples = torch.empty((0,3,64,32))\n",
    "for i in tqdm(inds):\n",
    "    i = int(i)\n",
    "    labs.append(testset[i]['lab'])\n",
    "    samples = torch.cat((samples,testset[i]['img'].unsqueeze(0)))\n",
    "labs = torch.tensor(labs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_resized = samples\n",
    "test_subset = torch.utils.data.TensorDataset(samples_resized,labs)\n",
    "testloader2 = torch.utils.data.DataLoader(test_subset, batch_size=128, shuffle=True, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial\n",
      "31.7391357421875 4.23046875 0.2156662940979004\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# ===================== Storing Batch Statistics =====================\n",
    "\n",
    "batch_info = {}\n",
    "train_embs_pkl_path = os.path.join(save_dir,f'{expt_name}.pkl')\n",
    "ckpt_path = os.path.join(save_dir,f'{expt_name}.pt')\n",
    "batch_info['readme'] = f'---- K={K} ---- ckpt_path={ckpt_path} ---- pkl_path={train_embs_pkl_path} ----'\n",
    "batch_info['imgs'] = [] # test images\n",
    "batch_info['preds'] = [] # corresponding model predictions\n",
    "batch_info['labels'] = [] # labels of the test images\n",
    "batch_info['pd'] = [] # corresponding prediction depths\n",
    "batch_info['layers_knn_prob'] = [] # for each test image we have a list of knn means for every layer\n",
    "batch_info['layers_knn_mode'] = [] # for each test image we have a list of knn mode for every layer\n",
    "\n",
    "print_memory_profile('Initial')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  0%|          | 0/20 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model output: tensor([[ 8.1206,  1.4472, -0.6958, -1.4942, -1.0209, -0.5914, -1.6754, -0.5160,\n",
      "         -1.9417, -1.6489]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.236328125 0.2197422981262207\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.6328125 8.52247667312622\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 294.26it/s]\n",
      "  5%|▌         | 1/20 [00:08<02:42,  8.53s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.5859375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 0\n",
      "knn_preds_mode: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "knn_preds_prob: [0.6206896305084229, 0.7586206793785095, 0.5517241358757019, 0.7586206793785095, 1.0, 0.7586206793785095, 0.8620689511299133, 0.9655172228813171, 1.0, 1.0, 0.9655172228813171, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]\n",
      "label: 0\n",
      "pred: 0\n",
      "\n",
      "\n",
      "Model output: tensor([[-11.1930,  60.9860,  -3.4734,  -7.9542,  -5.2608,  -2.8297,  -9.2390,\n",
      "          -2.1497, -10.3728,  -8.6765]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 295.39it/s]\n",
      " 10%|█         | 2/20 [00:17<02:35,  8.65s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 1\n",
      "knn_preds_mode: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\n",
      "knn_preds_prob: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]\n",
      "label: 1\n",
      "pred: 1\n",
      "\n",
      "\n",
      "Model output: tensor([[-21.2277,  99.6399,  -5.4229, -12.5495,  -8.2169,  -4.4475, -14.5344,\n",
      "          -3.4263, -16.3291, -13.6446]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 295.67it/s]\n",
      " 15%|█▌        | 3/20 [00:25<02:26,  8.63s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 2\n",
      "knn_preds_mode: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\n",
      "knn_preds_prob: [0.931034505367279, 0.931034505367279, 0.9655172228813171, 0.931034505367279, 1.0, 0.931034505367279, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]\n",
      "label: 1\n",
      "pred: 1\n",
      "\n",
      "\n",
      "Model output: tensor([[ 7.6705,  2.6307, -0.7396, -1.6109, -1.1043, -0.6339, -1.8148, -0.5479,\n",
      "         -2.0987, -1.7812]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 294.75it/s]\n",
      " 20%|██        | 4/20 [00:34<02:17,  8.61s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 3\n",
      "knn_preds_mode: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "knn_preds_prob: [1.0, 1.0, 0.9655172228813171, 1.0, 1.0, 1.0, 1.0, 1.0, 0.931034505367279, 0.9655172228813171, 0.931034505367279, 0.931034505367279, 0.8965517282485962, 0.931034505367279, 0.8275862336158752, 1.0, 1.0]\n",
      "label: 0\n",
      "pred: 0\n",
      "\n",
      "\n",
      "Model output: tensor([[ 9.0085,  3.1016, -0.8882, -1.9005, -1.3018, -0.7447, -2.1501, -0.6469,\n",
      "         -2.4603, -2.1067]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 296.25it/s]\n",
      " 25%|██▌       | 5/20 [00:43<02:09,  8.61s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 4\n",
      "knn_preds_mode: [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "knn_preds_prob: [0.5862069129943848, 0.5862069129943848, 0.5517241358757019, 0.5862069129943848, 0.7931034564971924, 0.6896551847457886, 0.8275862336158752, 0.9655172228813171, 0.9655172228813171, 0.931034505367279, 0.7931034564971924, 0.7931034564971924, 0.8620689511299133, 0.8620689511299133, 1.0, 1.0, 1.0]\n",
      "label: 0\n",
      "pred: 0\n",
      "\n",
      "\n",
      "Model output: tensor([[ 7.6680,  1.7886, -0.6971, -1.4793, -1.0156, -0.5833, -1.6722, -0.5049,\n",
      "         -1.9195, -1.6351]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 295.42it/s]\n",
      " 30%|███       | 6/20 [00:51<02:00,  8.61s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 5\n",
      "knn_preds_mode: [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "knn_preds_prob: [0.7586206793785095, 0.7586206793785095, 0.8275862336158752, 0.7931034564971924, 0.7586206793785095, 0.7586206793785095, 0.5517241358757019, 0.6206896305084229, 0.7586206793785095, 0.8275862336158752, 0.7931034564971924, 0.931034505367279, 0.9655172228813171, 0.9655172228813171, 1.0, 1.0, 1.0]\n",
      "label: 0\n",
      "pred: 0\n",
      "\n",
      "\n",
      "Model output: tensor([[ 7.3121,  2.2429, -0.7035, -1.5081, -1.0420, -0.5936, -1.6815, -0.5169,\n",
      "         -1.9383, -1.6669]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 295.38it/s]\n",
      " 35%|███▌      | 7/20 [01:00<01:51,  8.61s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 6\n",
      "knn_preds_mode: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "knn_preds_prob: [0.7931034564971924, 0.7586206793785095, 0.7241379022598267, 0.7931034564971924, 0.7931034564971924, 0.8275862336158752, 0.931034505367279, 0.8965517282485962, 0.931034505367279, 0.8965517282485962, 0.931034505367279, 0.8620689511299133, 0.9655172228813171, 0.931034505367279, 0.8965517282485962, 1.0, 1.0]\n",
      "label: 0\n",
      "pred: 0\n",
      "\n",
      "\n",
      "Model output: tensor([[ 7.8586,  1.3634, -0.6666, -1.4444, -0.9928, -0.5745, -1.6088, -0.5041,\n",
      "         -1.8695, -1.5944]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 295.86it/s]\n",
      " 40%|████      | 8/20 [01:08<01:43,  8.61s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 7\n",
      "knn_preds_mode: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "knn_preds_prob: [0.8965517282485962, 0.8965517282485962, 0.8965517282485962, 0.8620689511299133, 0.931034505367279, 0.7586206793785095, 0.8965517282485962, 0.931034505367279, 0.7931034564971924, 0.8275862336158752, 0.8965517282485962, 0.8620689511299133, 0.8965517282485962, 0.7931034564971924, 1.0, 1.0, 1.0]\n",
      "label: 0\n",
      "pred: 0\n",
      "\n",
      "\n",
      "Model output: tensor([[ 8.4332,  2.4870, -0.7930, -1.7028, -1.1679, -0.6679, -1.9333, -0.5767,\n",
      "         -2.2181, -1.8848]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 294.49it/s]\n",
      " 45%|████▌     | 9/20 [01:17<01:34,  8.61s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 8\n",
      "knn_preds_mode: [1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "knn_preds_prob: [0.5517241358757019, 0.6206896305084229, 0.6896551847457886, 0.5862069129943848, 0.931034505367279, 0.7241379022598267, 0.6551724076271057, 0.8620689511299133, 0.7931034564971924, 0.931034505367279, 0.931034505367279, 0.8965517282485962, 0.931034505367279, 0.931034505367279, 0.9655172228813171, 1.0, 1.0]\n",
      "label: 0\n",
      "pred: 0\n",
      "\n",
      "\n",
      "Model output: tensor([[-9.2904, 55.0931, -3.2007, -7.3103, -4.8395, -2.6093, -8.4798, -2.0040,\n",
      "         -9.5399, -7.9836]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 295.17it/s]\n",
      " 50%|█████     | 10/20 [01:26<01:26,  8.61s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 9\n",
      "knn_preds_mode: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\n",
      "knn_preds_prob: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]\n",
      "label: 1\n",
      "pred: 1\n",
      "\n",
      "\n",
      "Model output: tensor([[-21.8482, 101.5881,  -5.5150, -12.7686,  -8.3570,  -4.5232, -14.7975,\n",
      "          -3.4627, -16.6053, -13.8856]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 295.27it/s]\n",
      " 55%|█████▌    | 11/20 [01:34<01:17,  8.61s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 10\n",
      "knn_preds_mode: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\n",
      "knn_preds_prob: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]\n",
      "label: 1\n",
      "pred: 1\n",
      "\n",
      "\n",
      "Model output: tensor([[-21.4469, 100.9022,  -5.4952, -12.7013,  -8.3281,  -4.5108, -14.7266,\n",
      "          -3.4669, -16.5454, -13.8150]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 293.29it/s]\n",
      " 60%|██████    | 12/20 [01:43<01:08,  8.62s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 11\n",
      "knn_preds_mode: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\n",
      "knn_preds_prob: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]\n",
      "label: 1\n",
      "pred: 1\n",
      "\n",
      "\n",
      "Model output: tensor([[ 7.1601,  3.2328, -0.7531, -1.6389, -1.1236, -0.6429, -1.8380, -0.5517,\n",
      "         -2.1156, -1.8080]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 295.04it/s]\n",
      " 65%|██████▌   | 13/20 [01:51<01:00,  8.63s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 12\n",
      "knn_preds_mode: [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "knn_preds_prob: [0.517241358757019, 0.5862069129943848, 0.5517241358757019, 0.6206896305084229, 0.931034505367279, 0.6551724076271057, 0.8275862336158752, 0.9655172228813171, 1.0, 1.0, 1.0, 1.0, 1.0, 0.9655172228813171, 0.9655172228813171, 1.0, 1.0]\n",
      "label: 0\n",
      "pred: 0\n",
      "\n",
      "\n",
      "Model output: tensor([[-16.4019,  84.7092,  -4.7339, -10.9153,  -7.1605,  -3.8812, -12.6447,\n",
      "          -2.9879, -14.2104, -11.8654]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 293.31it/s]\n",
      " 70%|███████   | 14/20 [02:00<00:51,  8.64s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 13\n",
      "knn_preds_mode: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\n",
      "knn_preds_prob: [1.0, 1.0, 0.9655172228813171, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]\n",
      "label: 1\n",
      "pred: 1\n",
      "\n",
      "\n",
      "Model output: tensor([[ 8.3624,  1.4300, -0.7180, -1.5355, -1.0506, -0.6099, -1.7191, -0.5349,\n",
      "         -1.9813, -1.6948]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 294.35it/s]\n",
      " 75%|███████▌  | 15/20 [02:09<00:43,  8.64s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 14\n",
      "knn_preds_mode: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "knn_preds_prob: [0.6896551847457886, 0.6551724076271057, 0.5862069129943848, 0.6896551847457886, 0.8620689511299133, 0.8275862336158752, 0.9655172228813171, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]\n",
      "label: 0\n",
      "pred: 0\n",
      "\n",
      "\n",
      "Model output: tensor([[ 2.5750, 17.3441, -1.4298, -3.1656, -2.1530, -1.1738, -3.6579, -0.9378,\n",
      "         -4.1345, -3.4765]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 295.78it/s]\n",
      " 80%|████████  | 16/20 [02:17<00:34,  8.64s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 15\n",
      "knn_preds_mode: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\n",
      "knn_preds_prob: [0.8965517282485962, 0.8275862336158752, 0.931034505367279, 0.8620689511299133, 0.7931034564971924, 0.7586206793785095, 0.8275862336158752, 0.8275862336158752, 0.8275862336158752, 0.7931034564971924, 0.8965517282485962, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]\n",
      "label: 1\n",
      "pred: 1\n",
      "\n",
      "\n",
      "Model output: tensor([[-45.0261, 188.1541,  -9.8337, -22.8942, -14.9361,  -8.1209, -26.5865,\n",
      "          -6.2029, -29.8018, -24.8660]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 294.77it/s]\n",
      " 85%|████████▌ | 17/20 [02:26<00:25,  8.65s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 16\n",
      "knn_preds_mode: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\n",
      "knn_preds_prob: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]\n",
      "label: 1\n",
      "pred: 1\n",
      "\n",
      "\n",
      "Model output: tensor([[ 8.6258,  1.1382, -0.7080, -1.5219, -1.0402, -0.6021, -1.7064, -0.5284,\n",
      "         -1.9793, -1.6783]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 294.88it/s]\n",
      " 90%|█████████ | 18/20 [02:35<00:17,  8.69s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 17\n",
      "knn_preds_mode: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "knn_preds_prob: [0.9655172228813171, 1.0, 0.931034505367279, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]\n",
      "label: 0\n",
      "pred: 0\n",
      "\n",
      "\n",
      "Model output: tensor([[-22.4882, 103.1727,  -5.5762, -12.9175,  -8.4507,  -4.5719, -14.9605,\n",
      "          -3.5254, -16.8022, -14.0504]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 296.53it/s]\n",
      " 95%|█████████▌| 19/20 [02:44<00:08,  8.68s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 18\n",
      "knn_preds_mode: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\n",
      "knn_preds_prob: [0.9655172228813171, 0.8965517282485962, 0.9655172228813171, 0.931034505367279, 0.9655172228813171, 0.931034505367279, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]\n",
      "label: 1\n",
      "pred: 1\n",
      "\n",
      "\n",
      "Model output: tensor([[ 7.4415,  3.6513, -0.8054, -1.7361, -1.1922, -0.6770, -1.9792, -0.5750,\n",
      "         -2.2610, -1.9197]], device='cuda:0')\n",
      "Model forward pass\n",
      "31.7391357421875 4.361328125 0.34285879135131836\n",
      "\n",
      "\n",
      "Pickle load\n",
      "31.7391357421875 8.7578125 8.645593166351318\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "17it [00:00, 295.10it/s]\n",
      "100%|██████████| 20/20 [02:52<00:00,  8.63s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pickle batch processed\n",
      "31.7391357421875 10.7109375 8.522485256195068\n",
      "\n",
      "\n",
      "After GPU memory freed\n",
      "31.7391357421875 4.357421875 0.34279775619506836\n",
      "\n",
      "\n",
      "Test Image: 19\n",
      "knn_preds_mode: [1, 1, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "knn_preds_prob: [0.6896551847457886, 0.6896551847457886, 0.7241379022598267, 0.6551724076271057, 0.8275862336158752, 0.6551724076271057, 0.5862069129943848, 0.5517241358757019, 0.7931034564971924, 0.7586206793785095, 0.8620689511299133, 1.0, 0.8965517282485962, 0.9655172228813171, 0.931034505367279, 1.0, 0.9655172228813171]\n",
      "label: 0\n",
      "pred: 0\n",
      "\n",
      "\n",
      "Invalid Counts Ratio: 0/20\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# loop over test images\n",
    "invalid_counter = 0 # for invalid predictions (last layer mode != model output)\n",
    "for test_id, (img,lab) in enumerate(tqdm(test_subset)):\n",
    "\n",
    "    batch_info['imgs'].append(img)\n",
    "    with torch.no_grad():\n",
    "        to_pil_trans = transforms.ToPILImage()\n",
    "        img = to_pil_trans(img.to('cuda'))\n",
    "        img = trans(img).unsqueeze(0)\n",
    "        lab = int(lab)\n",
    "        if img.shape[1]==4:\n",
    "            img = img[:,0,:,:].unsqueeze(0)\n",
    "        feature_maps = []\n",
    "        out = net(img)\n",
    "        print(f'Model output: {out}')\n",
    "        batch_info['preds'].append(int(out.argmax()))\n",
    "        batch_info['labels'].append(lab)\n",
    "\n",
    "        print_memory_profile('Model forward pass')\n",
    "        with open(train_embs_pkl_path, 'rb') as handle:            \n",
    "            info_dict = pickle.load(handle)\n",
    "            print_memory_profile('Pickle load')\n",
    "\n",
    "            # loop over layers in densenet, and compute KNN for this test image\n",
    "            knn_preds_mode = []  # layer-wise final KNN classification preds   \n",
    "            knn_preds_prob = []\n",
    "            for layer_id,feat in tqdm(enumerate(feature_maps)):\n",
    "                X_i = feat.unsqueeze(1)  # test set\n",
    "                X_j = info_dict['feats'][layer_id].unsqueeze(0)  # train set\n",
    "                if lp_norm==2:\n",
    "                    D_ij = ((X_i - X_j) ** 2).sum(-1)  # symbolic matrix of squared L2 distances\n",
    "                elif lp_norm==1:\n",
    "                    D_ij = (abs(X_i - X_j)).sum(-1)  # symbolic matrix of squared L2 distances\n",
    "                else:\n",
    "                    raise('Invalid lp_norm in arguments!')\n",
    "\n",
    "                ind_knn = torch.topk(-D_ij,K,dim=1)  # Samples <-> Dataset, (N_test, K)\n",
    "                lab_knn = info_dict['labels'][ind_knn[1]]  # (N_test, K) array of integers in [0,9]\n",
    "                mode = int(lab_knn.squeeze().mode()[0])\n",
    "                knn_preds_mode.append(mode)\n",
    "                knn_preds_prob.append(float((lab_knn==mode).float().mean()))\n",
    "\n",
    "            print_memory_profile('Pickle batch processed')\n",
    "\n",
    "            # free GPU memory\n",
    "            del info_dict\n",
    "            torch.cuda.empty_cache()\n",
    "            print_memory_profile('After GPU memory freed') \n",
    "\n",
    "            print('Test Image: %d' %(test_id))\n",
    "            print(f'knn_preds_mode: {knn_preds_mode}')\n",
    "            print(f'knn_preds_prob: {knn_preds_prob}')\n",
    "            print(f'label: {lab}')\n",
    "            print(f'pred: {int(out.argmax())}')\n",
    "            print('\\n')\n",
    "            batch_info['layers_knn_prob'].append(knn_preds_prob)\n",
    "            batch_info['layers_knn_mode'].append(knn_preds_mode)\n",
    "            if int(out.argmax())==knn_preds_mode[-1]: # PD accurate\n",
    "                if knn_pos_thresh==0.5 and knn_neg_thresh==0.5:\n",
    "                    batch_info['pd'].append(compute_pred_depth(knn_preds_mode))\n",
    "                else:\n",
    "                    raise('Code not ready yet! Compute pred arr function also has to be updated!')\n",
    "            else: # PD inaccurate, KNN pred doesn't match model pred\n",
    "                print('Invalid datapoint: last_layer_mode != model_output')\n",
    "                invalid_counter += 1\n",
    "                batch_info['pd'].append(-99)\n",
    "print(f'Invalid Counts Ratio: {invalid_counter}/{num_test_imgs}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ===================== Save results =====================\n",
    "\n",
    "with open(os.path.join(save_dir,expt_name+'_testPDinfo.pkl'), 'wb') as handle:\n",
    "    pickle.dump(batch_info, handle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAb0AAAG5CAYAAAAav+pSAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8/fFQqAAAACXBIWXMAAAsTAAALEwEAmpwYAAAaa0lEQVR4nO3de5SkdX3n8fdHRhQQIkhjhIGMepSNy7qAE8WgUUFd1hAwG7NCAod42dl1RdQECYpnjbmcGGXVqInuHLkYJboKqOh6AY0GXZEw3BQZEVSEQS5NCKDICsh3/6hnQtPTl5qhq56u+b1f5/TpujxVz3f6TM97nqeeeipVhSRJLXhY3wNIkjQuRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPWkBSa5N8vy+59goyXOT3J/kpzO+jpljuV2STCf5+jzPc0ySSvLKGbftk+SLSW5NMucbeJMckWR9kruSfD/Js5fuTyeN3oq+B5C02X5cVSsXWeavgPXM8R/bJDsDbwS+M+uue4GPA38LfGqOx72ge96XAv8EPG5zB5f65paeNI8kHwb2Aj7TbVGdkOSAJN9IcnuSy5M8d8byX03y5939P03ymSSPSXJGkjuTXJRk1YzlK8lxSX7QbV29I8lD/p1M8kxgH+C0eRb5S+A9wK0zb6yqq6rqFDaN4UZvBf60qr5ZVfdX1Q1VdcNDnVcaJ6MnzaOqjgauA36rqh4FnAH8H+DPgV2A44GzkkzNeNgRwNHAHsATgQsYxGcXBlteb5m1mt8GVgP7A4cDLx9itN2S3Jzkh0nelWSHjXck2Qb4G+BYYJNdlEme3q3vA0OsZ+bjtukeN5XkmiQbkrwvyXab8zxS34yeNLyjgM9V1ee6LZ3zgHXAi2Ysc1pVfb+q7gA+D3y/qr5UVfcBnwD2m/Wcf1VVt1XVdcC7gSMXmeG7wL4Mdi0eBDwNeOeM+48DLqyqi2c/sAvX3wKvqar7h/oTP+CxwMOBlwDP7mbYD3jzZj6P1CujJw3vV4Df7XZt3p7kduBZPPi1rZtnXL57juuPmvWc18+4/CNg94UGqKqbqurKLro/BE5gECKS7M4geifN8/D/Dnyrqi5YaB3zuLv7/t6qurGqbmUQ2xct8Bhp2fFAFmlhM3cRXg98uKr+yxI+/5488BraXsCPN/PxBaS7/HQGAb4yCcB2wHZJbmKwu/Vg4DlJNoZqF2C/JPtW1bELrqTqX5JsYI5dptIkMXrSwm4GntBd/ghwUZL/AHyJwe6+A4BrqmrDFj7/G5JcyGAL8LU8eFflJroDZ37AIMArgbcBn+7u/jywasbiLwV+Dzi8qn6R5A+AR864/2zgTOCU7rkDPALYtrv+SKCq6ufd8qcBr0nyBQZHer4O+Ozm/XGlfrl7U1rYXwJv7nZlvpTBwSZvAqYZhOcNPLTfo08DFwOXMThI5pRFlt+fwcExdwHfAK5gsEuTqvp5t/vzpqq6CbgDuLe7TFXdPuv+e4A7u9cfYbD79m4e2PK8G7hqxrr/DLgI+B6Dg3IuBf5iS//gUh/ih8hK/ejeAP6kqrqm71mkVrilJ0lqhtGTlpkkH5h1mrGNX5v13jpJm3L3piSpGW7pSZKaMRFvWdh1111r1apVfY8hSVpGLr744luramrxJR8wEdFbtWoV69at63sMSdIykuRHm/sYd29Kkpph9CRJzTB6kqRmGD1JUjOMniSpGUZPktQMoydJaobRkyQ1w+hJkpph9CRJzTB6kqRmGD1JUjOMniSpGUZPktQMoydJaobRkyQ1YyI+RHZUjj/3+C1+7MkvPHkJJ5EkjYNbepKkZhg9SVIzjJ4kqRlGT5LUDKMnSWqG0ZMkNcPoSZKaYfQkSc0wepKkZhg9SVIzjJ4kqRlGT5LUDKMnSWqG0ZMkNcPoSZKaYfQkSc0wepKkZhg9SVIzjJ4kqRlGT5LUjJFFL8mpSW5JcsUc9x2fpJLsOqr1S5I02yi39E4HDpl9Y5I9gRcA141w3ZIkbWJk0auq84Hb5rjrXcAJQI1q3ZIkzWWsr+klOQy4oaouH+d6JUkCWDGuFSXZHjgJeOGQy68B1gDstddeI5xMktSKcW7pPRF4PHB5kmuBlcAlSX55roWram1Vra6q1VNTU2McU5K0tRrbll5VfRvYbeP1Lnyrq+rWcc0gSWrbKN+y8FHgAmDvJBuSvGJU65IkaRgj29KrqiMXuX/VqNYtSdJcPCOLJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPUlSM0YWvSSnJrklyRUzbntHku8m+VaSTyZ59KjWL0nSbKPc0jsdOGTWbecB+1TVU4HvAW8c4folSXqQkUWvqs4Hbpt127lVdV939ZvAylGtX5Kk2fp8Te/lwOfnuzPJmiTrkqybnp4e41iSpK1VL9FLchJwH3DGfMtU1dqqWl1Vq6empsY3nCRpq7Vi3CtMcgxwKHBwVdW41y9JatdYo5fkEOCPgedU1c/GuW5Jkkb5loWPAhcAeyfZkOQVwPuAHYHzklyW5AOjWr8kSbONbEuvqo6c4+ZTRrU+SZIW4xlZJEnNMHqSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDVjZNFLcmqSW5JcMeO2XZKcl+Tq7vvOo1q/JEmzjXJL73TgkFm3nQh8uaqeBHy5uy5J0liMLHpVdT5w26ybDwc+1F3+EPDiUa1fkqTZxv2a3mOr6kaA7vtu8y2YZE2SdUnWTU9Pj21ASdLWa9keyFJVa6tqdVWtnpqa6nscSdJWYNzRuznJ4wC677eMef2SpIaNO3rnAMd0l48BPj3m9UuSGjbKtyx8FLgA2DvJhiSvAN4GvCDJ1cALuuuSJI3FilE9cVUdOc9dB49qnZIkLWTZHsgiSdJSM3qSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnN2KzoJXlYkp1GNYwkSaO0aPSS/H2SnZLsAFwJXJXkDaMfTZKkpTXMlt5TqupO4MXA54C9gKNHOZQkSaMwTPQenuThDKL36aq6F6iRTiVJ0ggME73/BVwL7ACcn+RXgDtHOZQkSaOwYrEFquo9wHtm3PSjJM8b3UiSJI3GMAeyPDbJKUk+311/CnDMyCeTJGmJDbN783Tgi8Du3fXvAa8b0TySJI3MMNHbtao+DtwPUFX3Ab8Y6VSSJI3AMNG7K8lj6I7YTHIAcMdIp5IkaQQWPZAF+EPgHOCJSf4vMAW8ZKRTSZI0AsMcvXlJkucAewMBrureqydJ0kRZNHpJ/tOsm56c5A7g21V1y2jGkiRp6Q2ze/MVwDOBr3TXnwt8k0H8/rSqPjyi2SRJWlLDRO9+4Fer6mYYvG8PeD/wDOB8wOhJkibCMEdvrtoYvM4twJOr6jbA1/YkSRNjmC29ryX5LPCJ7vrvMDgH5w7A7aMaTJKkpTZM9F7NIHQHMjh68++As6qqAM/BKUmaGMO8ZaGAM7svSZIm1jAnnD4gyUVJfprkniS/SOJHC0mSJs4wB7K8DzgSuBrYDngl8N5RDiVJ0igM85oeVXVNkm2q6hfAaUm+MeK5JElacsNE72dJtgUuS/J24EYGn6K+xZK8nsEWYwHfBl5WVf/voTynJEmLGWb35tHANsCxwF3AngyO5twiSfYAjgNWV9U+3XMfsaXPJ0nSsIY5evNH3cW7gbcu4Xq3S3IvsD3w4yV6XkmS5jXM0ZuHJrk0yW1J7kzyk4dy9GZV3QCcDFzHYFfpHVV17pY+nyRJwxpm9+a7gWOAx1TVTlW1Y1XttKUrTLIzcDjweGB3YIckR82x3Jok65Ksm56e3tLVSZL0r4aJ3vXAFd2b1JfC84EfVtV097l8ZwO/PnuhqlpbVauravXU1NQSrVqS1LJhjt48Afhckn8Efr7xxqp65xau8zrggCTbM3id8GBg3RY+lyRJQxsmen8B/BR4JLDtQ11hVV2Y5EzgEuA+4FJg7UN9XkmSFjNM9Hapqhcu5Uqr6i3AW5byOSVJWswwr+l9KcmSRk+SpD4ME71XA19IcvdSvGVBkqS+DPPm9B3HMYgkSaM2b/SS7L/QA6vqkqUfR5Kk0VloS+9/LnBfAQct8SySJI3UvNGrqueNcxBJkkZtmANZJEnaKhg9SVIz5o1ekgO7748Y3ziSJI3OQlt67+m+XzCOQSRJGrWFjt68N8lpwB5J3jP7zqo6bnRjSZK09BaK3qEMPgboIODi8YwjSdLoLPSWhVuBjyVZX1WXj3EmSZJGYpijN/85ySeT3JLk5iRnJVk58skkSVpiw0TvNOAcYHdgD+Az3W2SJE2UYaK3W1WdVlX3dV+nA1MjnkuSpCU3TPSmkxyVZJvu6yjgn0c9mCRJS22Y6L0c+M/ATcCNwEu62yRJmijDfJ7edcBhY5hFkqSR8tybkqRmGD1JUjOMniSpGYtGL8mbZ1z2ExckSRNroY8WOiHJMxkcrbmRn7ggSZpYCx29eRXwu8ATknwNWA88JsneVXXVWKaTJGkJLbR781+ANwHXAM/lgc/XOzHJN0Y8lyRJS26hLb1DgLcATwTeCVwO3FVVLxvHYJIkLbV5t/Sq6k1VdTBwLfARBoGcSvL1JJ8Z03ySJC2ZRc/IAnyxqi4CLkryqqp6VpJdRz2YJElLbdG3LFTVCTOu/kF3262jGkiSpFHZrDen+wnqkqRJ5hlZJEnNMHqSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnN6CV6SR6d5Mwk302yvvuwWkmSRmqYE06Pwl8DX6iqlyTZFti+pzkkSQ0Ze/SS7AT8Bg+cvPoe4J5xzyFJak8fuzefAEwDpyW5NMkHk+wwe6Eka5KsS7Juenp6/FNKkrY6fURvBbA/8P6q2g+4Czhx9kJVtbaqVlfV6qmpqXHPKEnaCvURvQ3Ahqq6sLt+JoMISpI0UmOPXlXdBFyfZO/upoOBK8c9hySpPX0dvfka4IzuyM0fAC/raQ5JUkN6iV5VXQas7mPdkqR2eUYWSVIzjJ4kqRlGT5LUDKMnSWqG0ZMkNcPoSZKaYfQkSc0wepKkZhg9SVIzjJ4kqRlGT5LUDKMnSWqG0ZMkNcPoSZKaYfQkSc0wepKkZhg9SVIzjJ4kqRlGT5LUDKMnSWqG0ZMkNcPoSZKaYfQkSc0wepKkZhg9SVIzjJ4kqRlGT5LUDKMnSWqG0ZMkNcPoSZKaYfQkSc0wepKkZhg9SVIzjJ4kqRlGT5LUDKMnSWqG0ZMkNcPoSZKaYfQkSc3oLXpJtklyaZLP9jWDJKktfW7pvRZY3+P6JUmN6SV6SVYCvwl8sI/1S5La1NeW3ruBE4D751sgyZok65Ksm56eHttgkqSt19ijl+RQ4Jaqunih5apqbVWtrqrVU1NTY5pOkrQ162NL70DgsCTXAh8DDkrykR7mkCQ1ZuzRq6o3VtXKqloFHAH8Q1UdNe45JEnt8X16kqRmrOhz5VX1VeCrfc4gSWqHW3qSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPUlSM1b0PYDG5/hzj39Ijz/5hScv0SSS1A+39CRJzTB6kqRmGD1JUjOMniSpGUZPktQMoydJaobRkyQ1w+hJkpph9CRJzTB6kqRmGD1JUjOMniSpGUZPktQMoydJaobRkyQ1w+hJkpph9CRJzTB6kqRmGD1JUjOMniSpGWOPXpI9k3wlyfok30ny2nHPIElq04oe1nkf8EdVdUmSHYGLk5xXVVf2MIskqSFj39Krqhur6pLu8k+A9cAe455DktSePrb0/lWSVcB+wIVz3LcGWAOw1157jXcwCTj+3OMf0uNPfuHJSzSJpKXS24EsSR4FnAW8rqrunH1/Va2tqtVVtXpqamr8A0qStjq9RC/JwxkE74yqOruPGSRJ7enj6M0ApwDrq+qd416/JKldfWzpHQgcDRyU5LLu60U9zCFJaszYD2Spqq8DGfd6JUnyjCySpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnNMHqSpGYYPUlSM4yeJKkZRk+S1AyjJ0lqhtGTJDXD6EmSmmH0JEnN6CV6SQ5JclWSa5Kc2McMkqT2jD16SbYB/gb4j8BTgCOTPGXcc0iS2tPHlt7TgWuq6gdVdQ/wMeDwHuaQJDUmVTXeFSYvAQ6pqld2148GnlFVx85abg2wpru6D3DFWAfdcrsCt/Y9xGaYpHknaVaYrHmddXQmad5JmhVg76racXMesGJUkywgc9y2SXmrai2wFiDJuqpaPerBlsIkzQqTNe8kzQqTNa+zjs4kzTtJs8Jg3s19TB+7NzcAe864vhL4cQ9zSJIa00f0LgKelOTxSbYFjgDO6WEOSVJjxr57s6ruS3Is8EVgG+DUqvrOIg9bO/rJlswkzQqTNe8kzQqTNa+zjs4kzTtJs8IWzDv2A1kkSeqLZ2SRJDXD6EmSmrGsozdJpytLsmeSryRZn+Q7SV7b90yLSbJNkkuTfLbvWRaT5NFJzkzy3e5n/My+Z5pPktd3fweuSPLRJI/se6aZkpya5JYkV8y4bZck5yW5uvu+c58zbjTPrO/o/h58K8knkzy6xxEfZK55Z9x3fJJKsmsfs80236xJXtP9u/udJG/va76Z5vl7sG+Sbya5LMm6JE8f5rmWbfQm8HRl9wF/VFW/ChwAvHqZzwvwWmB930MM6a+BL1TVvwH+Pct07iR7AMcBq6tqHwYHax3R71SbOB04ZNZtJwJfrqonAV/uri8Hp7PprOcB+1TVU4HvAW8c91ALOJ1N5yXJnsALgOvGPdACTmfWrEmex+AMWU+tqn8LnNzDXHM5nU1/rm8H3lpV+wL/o7u+qGUbPSbsdGVVdWNVXdJd/gmDf5T36Heq+SVZCfwm8MG+Z1lMkp2A3wBOAaiqe6rq9l6HWtgKYLskK4DtWWbvQ62q84HbZt18OPCh7vKHgBePc6b5zDVrVZ1bVfd1V7/J4L2+y8I8P1uAdwEnMMeJOPoyz6yvAt5WVT/vlrll7IPNYZ5ZC9ipu/xLDPl7tpyjtwdw/YzrG1jGEZkpySpgP+DCnkdZyLsZ/BLe3/Mcw3gCMA2c1u2O/WCSHfoeai5VdQOD/x1fB9wI3FFV5/Y71VAeW1U3wuA/cMBuPc8zrJcDn+97iIUkOQy4oaou73uWITwZeHaSC5P8Y5Jf63ugBbwOeEeS6xn8zg21xb+cozfU6cqWmySPAs4CXldVd/Y9z1ySHArcUlUX9z3LkFYA+wPvr6r9gLtYPrvfHqR7Lexw4PHA7sAOSY7qd6qtU5KTGLyscEbfs8wnyfbASQx2v02CFcDODF6ieQPw8SRz/Vu8HLwKeH1V7Qm8nm5P0GKWc/Qm7nRlSR7OIHhnVNXZfc+zgAOBw5Jcy2C38UFJPtLvSAvaAGyoqo1bzmcyiOBy9Hzgh1U1XVX3AmcDv97zTMO4OcnjALrvy2K31nySHAMcCvx+Le83Gz+RwX+ALu9+31YClyT55V6nmt8G4Owa+CcGe4KWxYE3cziGwe8XwCcYvCS2qOUcvYk6XVn3v6FTgPVV9c6+51lIVb2xqlZW1SoGP9d/qKpluzVSVTcB1yfZu7vpYODKHkdayHXAAUm27/5OHMwyPehmlnMY/CNC9/3TPc6yoCSHAH8MHFZVP+t7noVU1berareqWtX9vm0A9u/+Ti9HnwIOAkjyZGBblu+nLvwYeE53+SDg6qEeVVXL9gt4EYOjs74PnNT3PIvM+iwGu1+/BVzWfb2o77mGmPu5wGf7nmOIOfcF1nU/308BO/c90wKzvhX4LoOPw/ow8Ii+Z5o130cZvN54L4N/hF8BPIbBUZtXd9936XvOBWa9hsHr/Rt/zz7Q95wLzTvr/muBXfuec4Gf7bbAR7q/u5cAB/U95wKzPgu4GLicwfETTxvmuTwNmSSpGct596YkSUvK6EmSmmH0JEnNMHqSpGYYPUlSM4ye1JMkP+17Bqk1Rk/aynUnvpaE0ZOWlSS/1Z3s99IkX0ry2CQP6z7nbqpb5mHdZ0zummQqyVlJLuq+DuyW+ZMka5OcC/xdr38oaRkxetLy8nXggBqcWPtjwAlVdT+Ds2T8frfM84HLq+pWBp8z+K6q+jXgd3jwR0U9DTi8qn5vbNNLy5y7PaTlZSXwv7uTPm8L/LC7/VQG58N8N4OP0zmtu/35wFNmnAh/pyQ7dpfPqaq7xzG0NCnc0pOWl/cC76uqfwf8V+CRAFV1PYNPQjgIeAYPfIbcw4BnVtW+3dceNfgQYxh8BJOkGYyetLz8EnBDd/mYWfd9kMFuzo9X1S+6284Fjt24QJJ9Rz2gNMmMntSf7ZNsmPH1h8CfAJ9I8jU2/UiXc4BH8cCuTYDjgNVJvpXkSuC/jWNwaVL5KQvShEiymsFBK8/uexZpUnkgizQBkpwIvIoHjuCUtAXc0pMkNcPX9CRJzTB6kqRmGD1JUjOMniSpGUZPktSM/w9qqrFacj0j7gAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 504x504 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# plot PDs histogram\n",
    "with open(os.path.join(save_dir,expt_name+'_testPDinfo.pkl'), 'rb') as handle:\n",
    "    batch_info = pickle.load(handle)\n",
    "\n",
    "batch_info['pd'] = np.array(batch_info['pd'])\n",
    "batch_info['labels'] = np.array(batch_info['labels'])\n",
    "batch_info['preds'] = np.array(batch_info['preds'])\n",
    "correct_preds_arr = (batch_info['preds']==batch_info['labels'])\n",
    "\n",
    "\n",
    "plt.figure(figsize=(7,7))\n",
    "plt.title(expt_name)\n",
    "plt.ylabel('# of Images')\n",
    "plt.xlabel('Layer')\n",
    "if model=='resnet18':\n",
    "    plt.xlim((0,18))\n",
    "elif model=='vgg16':\n",
    "    plt.xlim((0,16))\n",
    "# plt.ylim((0,80))\n",
    "plt.hist(batch_info['pd'][correct_preds_arr],bins=10,color='g',alpha=0.55)\n",
    "# plt.hist(batch_info['pd'][~correct_preds_arr ],bins=100,color='r',alpha=0.55)\n",
    "plt.savefig(os.path.join(save_dir,expt_name+'_PDplot.png'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([15.,  0.,  0.,  0.,  3.,  0.,  0.,  0.,  0.,  2.]),\n",
       " array([1. , 1.7, 2.4, 3.1, 3.8, 4.5, 5.2, 5.9, 6.6, 7.3, 8. ]),\n",
       " <BarContainer object of 10 artists>)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXAAAAD4CAYAAAD1jb0+AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8/fFQqAAAACXBIWXMAAAsTAAALEwEAmpwYAAANXklEQVR4nO3db4hl9X3H8fenjkF3o2jwJjWu2zElLIiUKEP+CRbcGDZRNA/6QMFgW2GetKmmLVbJg9AnxdCQptCSsqjREllpV0NDbFKXRLGCMZldNa6uadJko6smOyKNMS0Ym28fzBXWcWfu3XvO7L2/5v2CYeaee+aeL+vu2zPnnnMmVYUkqT2/Me0BJEmTMeCS1CgDLkmNMuCS1CgDLkmNmjueGzvjjDNqfn7+eG5Skpq3d+/eF6tqsHr5cQ34/Pw8S0tLx3OTktS8JD8+2nIPoUhSowy4JDXKgEtSowy4JDXKgEtSowy4JDXKgEtSowy4JDXKgEtSo47rlZhdzN9479S2ffDmS6e2bUlai3vgktQoAy5JjTLgktQoAy5JjTLgktQoAy5JjTLgktQoAy5JjTLgktSokQFPcluSw0n2H+W5P09SSc7YmPEkSWsZZw/8dmDH6oVJzgYuAZ7peSZJ0hhGBryqHgReOspTfwPcAFTfQ0mSRpvoGHiSy4HnqurxnueRJI3pmO9GmGQT8Cngw2OuvwgsAmzduvVYNydJWsMke+C/DZwDPJ7kILAF2JfkN4+2clXtrKqFqloYDAaTTypJeoNj3gOvqieAt7/+eBjxhap6sce5JEkjjHMa4S7gYWBbkkNJrt34sSRJo4zcA6+qq0Y8P9/bNJKksXklpiQ1yoBLUqMMuCQ1yoBLUqMMuCQ1yoBLUqMMuCQ1yoBLUqMMuCQ1yoBLUqMMuCQ1yoBLUqMMuCQ1yoBLUqMMuCQ1yoBLUqMMuCQ1yoBLUqMMuCQ1apxfanxbksNJ9h+x7K+TPJ3ku0m+nOS0DZ1SkvQm4+yB3w7sWLVsD3BeVf0O8B/ATT3PJUkaYWTAq+pB4KVVy+6rqteGD78FbNmA2SRJ6+jjGPgfAl9b68kki0mWkiwtLy/3sDlJEnQMeJJPAa8Bd661TlXtrKqFqloYDAZdNidJOsLcpN+Y5BrgMmB7VVV/I0mSxjFRwJPsAP4C+N2q+u9+R5IkjWOc0wh3AQ8D25IcSnIt8HfAKcCeJI8l+YcNnlOStMrIPfCquuooi2/dgFkkScfAKzElqVEGXJIaZcAlqVEGXJIaZcAlqVEGXJIaZcAlqVEGXJIaZcAlqVEGXJIaZcAlqVEGXJIaZcAlqVEGXJIaZcAlqVEGXJIaZcAlqVEGXJIaZcAlqVHj/FLj25IcTrL/iGVvS7InyfeHn0/f2DElSauNswd+O7Bj1bIbgW9U1buBbwwfS5KOo5EBr6oHgZdWLb4CuGP49R3Ax/odS5I0yqTHwN9RVS8ADD+/fa0VkywmWUqytLy8POHmJEmrbfibmFW1s6oWqmphMBhs9OYk6dfGpAH/aZIzAYafD/c3kiRpHJMG/CvANcOvrwH+pZ9xJEnjGuc0wl3Aw8C2JIeSXAvcDFyS5PvAJcPHkqTjaG7UClV11RpPbe95FknSMfBKTElqlAGXpEYZcElqlAGXpEYZcElqlAGXpEYZcElqlAGXpEYZcElqlAGXpEYZcElqlAGXpEYZcElqlAGXpEYZcElqlAGXpEYZcElqlAGXpEZ1CniSTyZ5Msn+JLuSnNTXYJKk9U0c8CRnAX8CLFTVecAJwJV9DSZJWl/XQyhzwMlJ5oBNwPPdR5IkjWPigFfVc8BngWeAF4CfVdV9q9dLsphkKcnS8vLy5JNKkt6gyyGU04ErgHOAdwKbk1y9er2q2llVC1W1MBgMJp9UkvQGXQ6hfAj4UVUtV9UvgXuAD/YzliRplC4BfwZ4f5JNSQJsBw70M5YkaZQux8AfAXYD+4Anhq+1s6e5JEkjzHX55qr6NPDpnmaRJB0Dr8SUpEYZcElqlAGXpEYZcElqlAGXpEYZcElqlAGXpEYZcElqlAGXpEYZcElqlAGXpEYZcElqlAGXpEYZcElqlAGXpEYZcElqlAGXpEYZcElqlAGXpEZ1CniS05LsTvJ0kgNJPtDXYJKk9XX6pcbA3wJfr6rfS/IWYFMPM0mSxjBxwJOcClwE/D5AVb0KvNrPWJKkUbocQnkXsAx8McmjSW5Jsnn1SkkWkywlWVpeXu6wOUnSkboEfA64APhCVZ0P/AK4cfVKVbWzqhaqamEwGHTYnCTpSF0Cfgg4VFWPDB/vZiXokqTjYOKAV9VPgGeTbBsu2g481ctUkqSRup6F8gngzuEZKD8E/qD7SJKkcXQKeFU9Biz0M4ok6Vh4JaYkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjOgc8yQlJHk3y1T4GkiSNp4898OuAAz28jiTpGHQKeJItwKXALf2MI0kaV9c98M8DNwC/6j6KJOlYTBzwJJcBh6tq74j1FpMsJVlaXl6edHOSpFW67IFfCFye5CBwF3Bxki+tXqmqdlbVQlUtDAaDDpuTJB1p4oBX1U1VtaWq5oErgW9W1dW9TSZJWpfngUtSo+b6eJGqegB4oI/XkiSNxz1wSWqUAZekRhlwSWqUAZekRhlwSWqUAZekRhlwSWqUAZekRvVyIY/Ul/kb753Kdg/efOlUtit14R64JDXKgEtSowy4JDXKgEtSowy4JDXKgEtSowy4JDXKgEtSowy4JDXKgEtSoyYOeJKzk9yf5ECSJ5Nc1+dgkqT1dbkXymvAn1XVviSnAHuT7Kmqp3qaTZK0jon3wKvqharaN/z658AB4Ky+BpMkra+XuxEmmQfOBx45ynOLwCLA1q1b+9icJE1kWne7hI2542XnNzGTvBW4G7i+ql5e/XxV7ayqhapaGAwGXTcnSRrqFPAkJ7IS7zur6p5+RpIkjaPLWSgBbgUOVNXn+htJkjSOLnvgFwIfBy5O8tjw46M9zSVJGmHiNzGr6iEgPc4iSToGXokpSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY3qFPAkO5J8L8kPktzY11CSpNEmDniSE4C/Bz4CnAtcleTcvgaTJK2vyx74e4EfVNUPq+pV4C7gin7GkiSNMtfhe88Cnj3i8SHgfatXSrIILA4fvpLkexNu7wzgxQm/t5N8ZqJvm9q8E2hpVtiAeSf8bzyOX/s/2w3U0qzkM53m/a2jLewS8BxlWb1pQdVOYGeH7axsLFmqqoWur3O8tDRvS7NCW/O2NCu0NW9Ls8LGzNvlEMoh4OwjHm8Bnu82jiRpXF0C/h3g3UnOSfIW4ErgK/2MJUkaZeJDKFX1WpI/Bv4NOAG4raqe7G2yN+t8GOY4a2nelmaFtuZtaVZoa96WZoUNmDdVbzpsLUlqgFdiSlKjDLgkNWrmA57ktiSHk+yf9iyjJDk7yf1JDiR5Msl1055pPUlOSvLtJI8P5/3Lac80SpITkjya5KvTnmWUJAeTPJHksSRL055nPUlOS7I7ydPDv78fmPZMa0mybfhn+vrHy0mun/Zca0nyyeG/r/1JdiU5qbfXnvVj4EkuAl4B/rGqzpv2POtJciZwZlXtS3IKsBf4WFU9NeXRjipJgM1V9UqSE4GHgOuq6ltTHm1NSf4UWABOrarLpj3PepIcBBaqauYvNklyB/DvVXXL8KyyTVX1X1Mea6ThLT2eA95XVT+e9jyrJTmLlX9X51bV/yT5J+Bfq+r2Pl5/5vfAq+pB4KVpzzGOqnqhqvYNv/45cICVK1ZnUq14ZfjwxOHHzP4fPckW4FLglmnP8v9JklOBi4BbAarq1RbiPbQd+M9ZjPcR5oCTk8wBm+jxepmZD3irkswD5wOPTHmUdQ0PSTwGHAb2VNUsz/t54AbgV1OeY1wF3Jdk7/CWErPqXcAy8MXh4albkmye9lBjuhLYNe0h1lJVzwGfBZ4BXgB+VlX39fX6BnwDJHkrcDdwfVW9PO151lNV/1tV72HlStr3JpnJw1RJLgMOV9Xeac9yDC6sqgtYuWPnHw0PB86iOeAC4AtVdT7wC2Dmbw89PNRzOfDP055lLUlOZ+Umf+cA7wQ2J7m6r9c34D0bHku+G7izqu6Z9jzjGv7I/ACwY7qTrOlC4PLhceW7gIuTfGm6I62vqp4ffj4MfJmVO3jOokPAoSN++trNStBn3UeAfVX102kPso4PAT+qquWq+iVwD/DBvl7cgPdo+KbgrcCBqvrctOcZJckgyWnDr09m5S/b01Mdag1VdVNVbamqeVZ+bP5mVfW2J9O3JJuHb2QzPBzxYWAmz6Sqqp8AzybZNly0HZjJN95XuYoZPnwy9Azw/iSbhn3Yzsp7Y72Y+YAn2QU8DGxLcijJtdOeaR0XAh9nZe/w9VOcPjrtodZxJnB/ku+ycm+bPVU186fnNeIdwENJHge+DdxbVV+f8kzr+QRw5/DvwnuAv5ruOOtLsgm4hJU92pk1/KlmN7APeIKV5vZ2Sf3Mn0YoSTq6md8DlyQdnQGXpEYZcElqlAGXpEYZcElqlAGXpEYZcElq1P8B79gV5UyjCjEAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.hist(batch_info['pd'][correct_preds_arr])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test Model on Particular Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 9/9 [00:37<00:00,  4.21s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "counter:0,total:180,accuracy:0.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "counter = 0\n",
    "total = 0\n",
    "for top in tqdm([0,1,2,4,5,6,7,8,9]):\n",
    "    for bot in [3]:\n",
    "        for i in range(20):\n",
    "            img_path = f'/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/data/domino_datasets/top_cifar10_bot_fmnist/train/top{top}_bot{bot}/{int(random.random()*1000)}.jpg'\n",
    "            ckpt_path = '/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/experiments/toy_expts/domino_expts/output/temp_54416.pt'\n",
    "\n",
    "            net = ResNet18(num_channels=num_ch)\n",
    "            net = net.to(device)\n",
    "            net.linear = nn.Linear(in_features=1024,out_features=10,bias=True)\n",
    "            net = torch.nn.DataParallel(net)\n",
    "            net.load_state_dict(torch.load(ckpt_path)['net'])\n",
    "            net.eval()\n",
    "\n",
    "            T = torchvision.transforms.Compose([\n",
    "                torchvision.transforms.ToTensor(),\n",
    "                torchvision.transforms.Resize((64,32)),\n",
    "            ])\n",
    "\n",
    "            with Image.open(img_path) as im:\n",
    "                out = net(T(im).unsqueeze(0))\n",
    "                total+=1\n",
    "                if int(out.argmax())==bot:\n",
    "                    counter+=1\n",
    "print(f'counter:{counter},total:{total},accuracy:{counter*100/total}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'counter:{counter},total:{total}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "0:48/180, 26%,21%,80%,4.5,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'temp_54416'"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Domino CSV Files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# user hyperparams\n",
    "top_dset_name = 'cifar10'\n",
    "bot_dset_name = 'fmnist'\n",
    "top_inds = [0,1]\n",
    "bot_inds = [0,1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "df = pd.DataFrame(columns=['path','bottom_label','top_label','val_train_split','test_split','all_zeros'])\n",
    "\n",
    "# domino = 5000/800\n",
    "# blank = 6000/1000\n",
    "if top_dset_name=='blank':\n",
    "    arr = np.arange(0,6000,1)\n",
    "    np.random.shuffle(arr)\n",
    "    val_arr = arr[:600]\n",
    "    train_arr = arr[600:]\n",
    "    num_test_imgs = 1000\n",
    "else:\n",
    "    arr = np.arange(0,5000,1)\n",
    "    np.random.shuffle(arr)\n",
    "    val_arr = arr[:500]\n",
    "    train_arr = arr[500:]\n",
    "    num_test_imgs = 800\n",
    "\n",
    "for (top_idx,bot_idx) in zip(top_inds,bot_inds):\n",
    "    for img_id in tqdm(train_arr):\n",
    "        if top_dset_name=='blank':\n",
    "            path = f'/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/data/domino_datasets/top_{top_dset_name}_bot_{bot_dset_name}/train/{bot_idx}/{img_id}.jpg'\n",
    "        else:    \n",
    "            path = f'/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/data/domino_datasets/top_{top_dset_name}_bot_{bot_dset_name}/train/top{top_idx}_bot{bot_idx}/{img_id}.jpg'\n",
    "        df = df.append({'path':path, 'bot_lbl':bot_idx, 'top_lbl':top_idx, 'val_train_split':1, 'test_split':0, 'all_zeros':0}, ignore_index=True)\n",
    "        \n",
    "    for img_id in tqdm(val_arr):\n",
    "        if top_dset_name=='blank':\n",
    "            path = f'/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/data/domino_datasets/top_{top_dset_name}_bot_{bot_dset_name}/train/{bot_idx}/{img_id}.jpg'\n",
    "        else:\n",
    "            path = f'/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/data/domino_datasets/top_{top_dset_name}_bot_{bot_dset_name}/train/top{top_idx}_bot{bot_idx}/{img_id}.jpg'\n",
    "        df = df.append({'path':path, 'bot_lbl':bot_idx, 'top_lbl':top_idx, 'val_train_split':0, 'test_split':0, 'all_zeros':0}, ignore_index=True)\n",
    "        \n",
    "    for img_id in tqdm(range(num_test_imgs)):\n",
    "        if top_dset_name=='blank':            \n",
    "            path = f'/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/data/domino_datasets/top_{top_dset_name}_bot_{bot_dset_name}/test/{bot_idx}/{img_id}.jpg'\n",
    "        else:\n",
    "            path = f'/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/data/domino_datasets/top_{top_dset_name}_bot_{bot_dset_name}/test/top{top_idx}_bot{bot_idx}/{img_id}.jpg'\n",
    "        df = df.append({'path':path, 'bot_lbl':bot_idx, 'top_lbl':top_idx, 'val_train_split':2, 'test_split':1, 'all_zeros':0}, ignore_index=True)\n",
    "     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# csv_save_path = '/jet/home/nmurali/asc170022p/nmurali/projects/shortcut_detection_and_mitigation/experiments/toy_expts/domino_expts/top_cifar10_bot_fmnist_2class_ro1p0.csv'\n",
    "# df.to_csv(csv_save_path, index=False)   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pl",
   "language": "python",
   "name": "pl"
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}