Train_baseline.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "aaee65d4-fdcd-492e-8a12-24a9f6e6e9af",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Goal : To classify the type of ovarian cancer from microscopy scans of biopsy samples.\n",
    "\n",
    "# Search SoTA model on Kaggle, Papers with code - Medical Data, Graph Transformer, GNN, Global and Local Analysis\n",
    "\n",
    "# You should manage, including differences in image dimensions, quality, slide staining techniques, and more\n",
    "\n",
    "# There are 2 types of images - TMA and WSI (25 images are TMA and 513(rest) are WSI)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9d1971b2-09c3-4a8d-b158-355c9772dfa9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/projectnb/cs640grp/students/avarshn/.conda/envs/ocean'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import sys\n",
    "sys.prefix"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48222a33-507b-4f98-8df3-8538ed479f1c",
   "metadata": {},
   "source": [
    "# Check if GPU is available"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "143df7b3-2bc3-4651-b9a8-ed3f5a8108be",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c193c329-3591-4572-9640-35b274bf754e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3c04ae72-d26a-4c3c-8415-ba36cb557c94",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/bin/bash: nvidia: command not found\n"
     ]
    }
   ],
   "source": [
    "!nvidia smi\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "150de13e-2b4e-499e-9092-e541baa686b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "02:00.0 3D controller: NVIDIA Corporation GP100GL [Tesla P100 PCIe 12GB] (rev a1)\n",
      "82:00.0 3D controller: NVIDIA Corporation GP100GL [Tesla P100 PCIe 12GB] (rev a1)\n"
     ]
    }
   ],
   "source": [
    "!lspci | grep -i nvidia"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75827adc-9618-4d47-8fdc-44dbdc43cbe8",
   "metadata": {},
   "source": [
    "# Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b7db4971-169f-47a7-9938-cc2d20e23b09",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from time import time\n",
    "from collections import defaultdict\n",
    "from tqdm import tqdm\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from sklearn.utils.class_weight import compute_class_weight\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import accuracy_score, f1_score, confusion_matrix\n",
    "\n",
    "from PIL import Image\n",
    "import cv2\n",
    "\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.optim import AdamW\n",
    "\n",
    "from torchsummary import summary\n",
    "\n",
    "import torchvision\n",
    "from torchvision import transforms, utils\n",
    "from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5c70d5b-6ba8-4faf-a3ff-5e37b6a0c467",
   "metadata": {},
   "source": [
    "# Configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4dcff8b2-936b-458c-a5f1-b7b711030b2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "CONFIG = {\n",
    "    \"seed\": 42,\n",
    "    # \"img_size\": 512,\n",
    "    # \"model_name\": \"tf_efficientnetv2_s_in21ft1k\",\n",
    "    \"num_classes\": 5,\n",
    "    # \"valid_batch_size\": 32,\n",
    "    \"batch_size\" : 8,\n",
    "    \"epochs\" : 10,\n",
    "    \"device\": torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\"),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "59ef42e6-00e5-40ff-a2b9-9e44b718dcee",
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_seed(seed=42):\n",
    "    '''Sets the seed of the entire notebook so results are the same every time we run.\n",
    "    This is for REPRODUCIBILITY.'''\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    # When running on the CuDNN backend, two further options must be set\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "    # Set a fixed value for the hash seed\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    \n",
    "set_seed(CONFIG['seed'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a226b920-d5c8-4257-8247-12ce684c5d67",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "f46a3116-7616-4ada-bfb5-fc823cfa1eb9",
   "metadata": {},
   "source": [
    "# Loading Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7e237727-16ea-461b-b540-5f0dee625bd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "SCC_folder = \"/projectnb/cs640grp/\"\n",
    "data_folder = \"materials/UBC-OCEAN_CS640/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6dddfe67-3b2f-40c9-b9a4-3914e670d3ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = SCC_folder + data_folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9099497d-6b7c-4993-bc95-aadde2dc98a8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['test_script.py',\n",
       " 'test_images_compressed_80',\n",
       " 'train.csv',\n",
       " 'test.csv',\n",
       " 'train_images_compressed_80',\n",
       " 'all_labels.npy']"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.listdir(data_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "683874df-8c8c-4e86-b543-4f12d6e9850f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['HGSC', 'EC', 'MC', 'CC', 'LGSC'], dtype='<U4')"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels = np.load(data_path +  'all_labels.npy')\n",
    "labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "efaaacb5-8ffa-492f-b87e-a0b5f57cde5f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>image_id</th>\n",
       "      <th>label</th>\n",
       "      <th>image_path</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>57598</td>\n",
       "      <td>CC</td>\n",
       "      <td>57598.jpg</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>30868</td>\n",
       "      <td>MC</td>\n",
       "      <td>30868.jpg</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>42549</td>\n",
       "      <td>CC</td>\n",
       "      <td>42549.jpg</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>64824</td>\n",
       "      <td>CC</td>\n",
       "      <td>64824.jpg</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>15293</td>\n",
       "      <td>HGSC</td>\n",
       "      <td>15293.jpg</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   image_id label image_path\n",
       "0     57598    CC  57598.jpg\n",
       "1     30868    MC  30868.jpg\n",
       "2     42549    CC  42549.jpg\n",
       "3     64824    CC  64824.jpg\n",
       "4     15293  HGSC  15293.jpg"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df = pd.read_csv(data_path +  'train.csv')\n",
    "train_df[\"image_path\"] = train_df[\"image_id\"].apply(lambda x : str(x) + \".jpg\")\n",
    "train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "57c0b006-58ba-4eec-968c-af28a91cdeed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 430 entries, 0 to 429\n",
      "Data columns (total 3 columns):\n",
      " #   Column      Non-Null Count  Dtype \n",
      "---  ------      --------------  ----- \n",
      " 0   image_id    430 non-null    int64 \n",
      " 1   label       430 non-null    object\n",
      " 2   image_path  430 non-null    object\n",
      "dtypes: int64(1), object(2)\n",
      "memory usage: 10.2+ KB\n"
     ]
    }
   ],
   "source": [
    "train_df.info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "f2ba11c1-7988-4d55-99ea-0ffa3f52ba31",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "430"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df[\"image_id\"].nunique()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b757d26-8798-43ad-a4d8-29782a89c84e",
   "metadata": {},
   "source": [
    "## Training Data Distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "10b86d40-0a58-4e00-a56d-d64311d63a0f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "label\n",
       "HGSC    177\n",
       "EC       99\n",
       "CC       79\n",
       "LGSC     38\n",
       "MC       37\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df[\"label\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d8cc3809-4a20-4554-abeb-ec115a63dcb6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "label\n",
       "HGSC    41.162791\n",
       "EC      23.023256\n",
       "CC      18.372093\n",
       "LGSC     8.837209\n",
       "MC       8.604651\n",
       "Name: proportion, dtype: float64"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "100*train_df[\"label\"].value_counts(normalize = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "0f6fa3d6-96cc-4e1b-b099-5c4b0fa7f315",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>image_id</th>\n",
       "      <th>label</th>\n",
       "      <th>image_path</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>7482</td>\n",
       "      <td>CC</td>\n",
       "      <td>7482.jpg</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   image_id label image_path\n",
       "0      7482    CC   7482.jpg"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_df = pd.read_csv(data_path +  'test.csv')\n",
    "test_df[\"image_path\"] = test_df[\"image_id\"].apply(lambda x : str(x) + \".jpg\")\n",
    "test_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "40d4f669-2d6c-4e56-b618-57311562cabc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# There is class imbalance\n",
    "\n",
    "# high-grade serous carcinoma, clear-cell ovarian carcinoma, endometrioid, low-grade serous, and mucinous carcinoma\n",
    "\n",
    "# Can use data augmentation to resolve this as slide can be in any fashion so can rotate the slides"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77e53b39-0440-4ad3-aa2e-b1ae00743088",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20f7d458-d0ad-4426-947e-1cefda8370fe",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "2f3350a3-92d5-4526-8c18-60568057476c",
   "metadata": {},
   "source": [
    "# Custom Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "02dd6b7a-b315-49e5-ab7d-50141f6a92c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "Image.MAX_IMAGE_PIXELS = None\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4e97aba-1582-4797-8187-9b2c1a82732c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "c442f33a-4624-4992-a37f-2be8d5806e25",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_weights = EfficientNet_B0_Weights.DEFAULT\n",
    "preprocess = model_weights.transforms()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "467b1f55-b77b-414d-9fa2-65ceaed0ddd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# preprocess = transforms.Compose([\n",
    "#         # transforms.RandomSizedCrop(224),\n",
    "#         # transforms.RandomHorizontalFlip(),\n",
    "#         transforms.Resize((4000, 4000)),\n",
    "#         transforms.ToTensor(),   # Converts a PIL Image or a NumPy array with values in the range [0, 255] to a PyTorch tensor\n",
    "#                                  # with values in the range [0.0, 1.0].\n",
    "#         # transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "#         #                      std=[0.229, 0.224, 0.225])\n",
    "# ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "002b334d-0f52-4bde-828c-eddcbb34a033",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ocean_dataset(Dataset):\n",
    "    \n",
    "    def __init__(self, dataframe, image_dir, transform = None):\n",
    "        \n",
    "        self.ocean_df = dataframe\n",
    "        self.ocean_df[\"image_path\"] = self.ocean_df[\"image_id\"].apply(lambda x : str(x) + \".jpg\")\n",
    "        \n",
    "        self.image_dir = image_dir\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.ocean_df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist()\n",
    "\n",
    "        img_path = os.path.join(self.image_dir, self.ocean_df[\"image_path\"][idx])\n",
    "        image = Image.open(img_path)\n",
    "\n",
    "        # # Convert Pillow image to NumPy array\n",
    "        # img_array = np.array(image)\n",
    "        \n",
    "        # # Display the range of pixel values\n",
    "        # print(f\"Minimum pixel value: {np.min(img_array)}\")\n",
    "        # print(f\"Maximum pixel value: {np.max(img_array)}\")\n",
    "\n",
    "\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "\n",
    "        label_mapper = {'HGSC' : 0, 'EC' : 1, 'MC' : 2, 'CC' : 3, 'LGSC' : 4}\n",
    "        label = label_mapper[self.ocean_df.iloc[idx, 1]]\n",
    "\n",
    "        \n",
    "        # sample = {'image': image, 'label': label}\n",
    "\n",
    "        return image, label\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33858d05-4c7f-4766-8cb5-84cc5b1dc770",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "40838c7b-7923-4e1e-b6e2-6a15d330293f",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(data_path +  'train.csv')\n",
    "\n",
    "# For testing the pipeline\n",
    "# df = df[0:30]\n",
    "\n",
    "train_df, test_df = train_test_split(df, test_size = 0.3, random_state = CONFIG['seed'])\n",
    "\n",
    "train_df = train_df.reset_index(drop=True)\n",
    "test_df = test_df.reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01054b83-dfda-4717-a6b5-64608e69e9d5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "27eba013-ad6c-4dc8-b686-e7e327206c69",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = ocean_dataset(train_df, data_path + 'train_images_compressed_80/', preprocess)\n",
    "\n",
    "test_dataset = ocean_dataset(test_df, data_path + 'train_images_compressed_80/', preprocess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "86204af9-161a-488c-b890-b8f11b17911b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# a,b = train_dataset.__getitem__(torch.tensor(8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "bd5247e2-7402-4bdd-b77d-f5d75b2ce125",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b46cb26c-214f-4d06-a8c4-905d0ec08154",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "20401933-47b3-404e-ba4e-be65e40f6ae6",
   "metadata": {},
   "source": [
    "# Class Weights - From Training Data since the dataset is imbalanced"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "31f40cfe-c181-440a-8f3e-d440782b30c3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0        CC\n",
       "1      HGSC\n",
       "2      HGSC\n",
       "3        MC\n",
       "4        MC\n",
       "       ... \n",
       "296    HGSC\n",
       "297    HGSC\n",
       "298      EC\n",
       "299      EC\n",
       "300    HGSC\n",
       "Name: label, Length: 301, dtype: object"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df[\"label\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "6e8ff4c7-62e4-4654-a53b-209323c92e3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "label_mapper = {'HGSC' : 0, 'EC' : 1, 'MC' : 2, 'CC' : 3, 'LGSC' : 4}\n",
    "labels = train_df[\"label\"].apply(lambda x: label_mapper[x])\n",
    "labels_arr = labels.values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "75b28ad6-975a-4fc2-a19b-73bde1c09870",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "label\n",
       "HGSC    115\n",
       "EC       73\n",
       "CC       63\n",
       "MC       27\n",
       "LGSC     23\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df[\"label\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "c8fbc531-f6d7-4203-8ed1-1a1ecad02f91",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Formula -\n",
    "# class_weight_of_that_class = n_samples / (n_classes * n_samples_of_that_class)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "d3548690-c5de-462c-b869-e2ed6836beb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# NumPy array of class labels - labels_arr\n",
    "\n",
    "class_weights = compute_class_weight('balanced',classes = np.unique(labels_arr), y = labels_arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "42486cca-117d-4dcb-9a8b-a2a118939821",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.52347826, 0.82465753, 2.22962963, 0.95555556, 2.6173913 ])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "4b19a847-2296-4df0-9897-8e5a295e6c6b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.5235, 0.8247, 2.2296, 0.9556, 2.6174], dtype=torch.float64)"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class_weights_tensor = torch.tensor(class_weights)\n",
    "class_weights_tensor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c34038c-0762-4061-85ee-cafd44eddddf",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "195ddb37-8f25-4c89-b174-8ca89565eeb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_model = efficientnet_b0(weights = model_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "e68d3d58-8c83-4095-975a-62d6105dacf7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 Sequential(\n",
      "  (0): Conv2dNormActivation(\n",
      "    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "    (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (2): SiLU(inplace=True)\n",
      "  )\n",
      "  (1): Sequential(\n",
      "    (0): MBConv(\n",
      "      (block): Sequential(\n",
      "        (0): Conv2dNormActivation(\n",
      "          (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
      "          (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (1): SqueezeExcitation(\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
      "          (fc1): Conv2d(32, 8, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (fc2): Conv2d(8, 32, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (activation): SiLU(inplace=True)\n",
      "          (scale_activation): Sigmoid()\n",
      "        )\n",
      "        (2): Conv2dNormActivation(\n",
      "          (0): Conv2d(32, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (stochastic_depth): StochasticDepth(p=0.0, mode=row)\n",
      "    )\n",
      "  )\n",
      "  (2): Sequential(\n",
      "    (0): MBConv(\n",
      "      (block): Sequential(\n",
      "        (0): Conv2dNormActivation(\n",
      "          (0): Conv2d(16, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (1): Conv2dNormActivation(\n",
      "          (0): Conv2d(96, 96, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=96, bias=False)\n",
      "          (1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (2): SqueezeExcitation(\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
      "          (fc1): Conv2d(96, 4, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (fc2): Conv2d(4, 96, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (activation): SiLU(inplace=True)\n",
      "          (scale_activation): Sigmoid()\n",
      "        )\n",
      "        (3): Conv2dNormActivation(\n",
      "          (0): Conv2d(96, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (stochastic_depth): StochasticDepth(p=0.0125, mode=row)\n",
      "    )\n",
      "    (1): MBConv(\n",
      "      (block): Sequential(\n",
      "        (0): Conv2dNormActivation(\n",
      "          (0): Conv2d(24, 144, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (1): Conv2dNormActivation(\n",
      "          (0): Conv2d(144, 144, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=144, bias=False)\n",
      "          (1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (2): SqueezeExcitation(\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
      "          (fc1): Conv2d(144, 6, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (fc2): Conv2d(6, 144, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (activation): SiLU(inplace=True)\n",
      "          (scale_activation): Sigmoid()\n",
      "        )\n",
      "        (3): Conv2dNormActivation(\n",
      "          (0): Conv2d(144, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (stochastic_depth): StochasticDepth(p=0.025, mode=row)\n",
      "    )\n",
      "  )\n",
      "  (3): Sequential(\n",
      "    (0): MBConv(\n",
      "      (block): Sequential(\n",
      "        (0): Conv2dNormActivation(\n",
      "          (0): Conv2d(24, 144, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (1): Conv2dNormActivation(\n",
      "          (0): Conv2d(144, 144, kernel_size=(5, 5), stride=(2, 2), padding=(2, 2), groups=144, bias=False)\n",
      "          (1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (2): SqueezeExcitation(\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
      "          (fc1): Conv2d(144, 6, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (fc2): Conv2d(6, 144, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (activation): SiLU(inplace=True)\n",
      "          (scale_activation): Sigmoid()\n",
      "        )\n",
      "        (3): Conv2dNormActivation(\n",
      "          (0): Conv2d(144, 40, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (stochastic_depth): StochasticDepth(p=0.037500000000000006, mode=row)\n",
      "    )\n",
      "    (1): MBConv(\n",
      "      (block): Sequential(\n",
      "        (0): Conv2dNormActivation(\n",
      "          (0): Conv2d(40, 240, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(240, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (1): Conv2dNormActivation(\n",
      "          (0): Conv2d(240, 240, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), groups=240, bias=False)\n",
      "          (1): BatchNorm2d(240, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (2): SqueezeExcitation(\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
      "          (fc1): Conv2d(240, 10, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (fc2): Conv2d(10, 240, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (activation): SiLU(inplace=True)\n",
      "          (scale_activation): Sigmoid()\n",
      "        )\n",
      "        (3): Conv2dNormActivation(\n",
      "          (0): Conv2d(240, 40, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (stochastic_depth): StochasticDepth(p=0.05, mode=row)\n",
      "    )\n",
      "  )\n",
      "  (4): Sequential(\n",
      "    (0): MBConv(\n",
      "      (block): Sequential(\n",
      "        (0): Conv2dNormActivation(\n",
      "          (0): Conv2d(40, 240, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(240, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (1): Conv2dNormActivation(\n",
      "          (0): Conv2d(240, 240, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=240, bias=False)\n",
      "          (1): BatchNorm2d(240, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (2): SqueezeExcitation(\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
      "          (fc1): Conv2d(240, 10, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (fc2): Conv2d(10, 240, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (activation): SiLU(inplace=True)\n",
      "          (scale_activation): Sigmoid()\n",
      "        )\n",
      "        (3): Conv2dNormActivation(\n",
      "          (0): Conv2d(240, 80, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(80, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (stochastic_depth): StochasticDepth(p=0.0625, mode=row)\n",
      "    )\n",
      "    (1): MBConv(\n",
      "      (block): Sequential(\n",
      "        (0): Conv2dNormActivation(\n",
      "          (0): Conv2d(80, 480, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(480, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (1): Conv2dNormActivation(\n",
      "          (0): Conv2d(480, 480, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=480, bias=False)\n",
      "          (1): BatchNorm2d(480, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (2): SqueezeExcitation(\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
      "          (fc1): Conv2d(480, 20, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (fc2): Conv2d(20, 480, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (activation): SiLU(inplace=True)\n",
      "          (scale_activation): Sigmoid()\n",
      "        )\n",
      "        (3): Conv2dNormActivation(\n",
      "          (0): Conv2d(480, 80, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(80, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (stochastic_depth): StochasticDepth(p=0.07500000000000001, mode=row)\n",
      "    )\n",
      "    (2): MBConv(\n",
      "      (block): Sequential(\n",
      "        (0): Conv2dNormActivation(\n",
      "          (0): Conv2d(80, 480, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(480, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (1): Conv2dNormActivation(\n",
      "          (0): Conv2d(480, 480, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=480, bias=False)\n",
      "          (1): BatchNorm2d(480, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (2): SqueezeExcitation(\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
      "          (fc1): Conv2d(480, 20, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (fc2): Conv2d(20, 480, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (activation): SiLU(inplace=True)\n",
      "          (scale_activation): Sigmoid()\n",
      "        )\n",
      "        (3): Conv2dNormActivation(\n",
      "          (0): Conv2d(480, 80, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(80, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (stochastic_depth): StochasticDepth(p=0.08750000000000001, mode=row)\n",
      "    )\n",
      "  )\n",
      "  (5): Sequential(\n",
      "    (0): MBConv(\n",
      "      (block): Sequential(\n",
      "        (0): Conv2dNormActivation(\n",
      "          (0): Conv2d(80, 480, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(480, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (1): Conv2dNormActivation(\n",
      "          (0): Conv2d(480, 480, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), groups=480, bias=False)\n",
      "          (1): BatchNorm2d(480, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (2): SqueezeExcitation(\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
      "          (fc1): Conv2d(480, 20, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (fc2): Conv2d(20, 480, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (activation): SiLU(inplace=True)\n",
      "          (scale_activation): Sigmoid()\n",
      "        )\n",
      "        (3): Conv2dNormActivation(\n",
      "          (0): Conv2d(480, 112, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(112, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (stochastic_depth): StochasticDepth(p=0.1, mode=row)\n",
      "    )\n",
      "    (1): MBConv(\n",
      "      (block): Sequential(\n",
      "        (0): Conv2dNormActivation(\n",
      "          (0): Conv2d(112, 672, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(672, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (1): Conv2dNormActivation(\n",
      "          (0): Conv2d(672, 672, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), groups=672, bias=False)\n",
      "          (1): BatchNorm2d(672, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (2): SqueezeExcitation(\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
      "          (fc1): Conv2d(672, 28, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (fc2): Conv2d(28, 672, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (activation): SiLU(inplace=True)\n",
      "          (scale_activation): Sigmoid()\n",
      "        )\n",
      "        (3): Conv2dNormActivation(\n",
      "          (0): Conv2d(672, 112, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(112, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (stochastic_depth): StochasticDepth(p=0.1125, mode=row)\n",
      "    )\n",
      "    (2): MBConv(\n",
      "      (block): Sequential(\n",
      "        (0): Conv2dNormActivation(\n",
      "          (0): Conv2d(112, 672, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(672, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (1): Conv2dNormActivation(\n",
      "          (0): Conv2d(672, 672, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), groups=672, bias=False)\n",
      "          (1): BatchNorm2d(672, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (2): SqueezeExcitation(\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
      "          (fc1): Conv2d(672, 28, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (fc2): Conv2d(28, 672, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (activation): SiLU(inplace=True)\n",
      "          (scale_activation): Sigmoid()\n",
      "        )\n",
      "        (3): Conv2dNormActivation(\n",
      "          (0): Conv2d(672, 112, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(112, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (stochastic_depth): StochasticDepth(p=0.125, mode=row)\n",
      "    )\n",
      "  )\n",
      "  (6): Sequential(\n",
      "    (0): MBConv(\n",
      "      (block): Sequential(\n",
      "        (0): Conv2dNormActivation(\n",
      "          (0): Conv2d(112, 672, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(672, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (1): Conv2dNormActivation(\n",
      "          (0): Conv2d(672, 672, kernel_size=(5, 5), stride=(2, 2), padding=(2, 2), groups=672, bias=False)\n",
      "          (1): BatchNorm2d(672, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (2): SqueezeExcitation(\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
      "          (fc1): Conv2d(672, 28, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (fc2): Conv2d(28, 672, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (activation): SiLU(inplace=True)\n",
      "          (scale_activation): Sigmoid()\n",
      "        )\n",
      "        (3): Conv2dNormActivation(\n",
      "          (0): Conv2d(672, 192, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (stochastic_depth): StochasticDepth(p=0.1375, mode=row)\n",
      "    )\n",
      "    (1): MBConv(\n",
      "      (block): Sequential(\n",
      "        (0): Conv2dNormActivation(\n",
      "          (0): Conv2d(192, 1152, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(1152, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (1): Conv2dNormActivation(\n",
      "          (0): Conv2d(1152, 1152, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), groups=1152, bias=False)\n",
      "          (1): BatchNorm2d(1152, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (2): SqueezeExcitation(\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
      "          (fc1): Conv2d(1152, 48, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (fc2): Conv2d(48, 1152, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (activation): SiLU(inplace=True)\n",
      "          (scale_activation): Sigmoid()\n",
      "        )\n",
      "        (3): Conv2dNormActivation(\n",
      "          (0): Conv2d(1152, 192, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (stochastic_depth): StochasticDepth(p=0.15000000000000002, mode=row)\n",
      "    )\n",
      "    (2): MBConv(\n",
      "      (block): Sequential(\n",
      "        (0): Conv2dNormActivation(\n",
      "          (0): Conv2d(192, 1152, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(1152, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (1): Conv2dNormActivation(\n",
      "          (0): Conv2d(1152, 1152, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), groups=1152, bias=False)\n",
      "          (1): BatchNorm2d(1152, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (2): SqueezeExcitation(\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
      "          (fc1): Conv2d(1152, 48, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (fc2): Conv2d(48, 1152, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (activation): SiLU(inplace=True)\n",
      "          (scale_activation): Sigmoid()\n",
      "        )\n",
      "        (3): Conv2dNormActivation(\n",
      "          (0): Conv2d(1152, 192, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (stochastic_depth): StochasticDepth(p=0.1625, mode=row)\n",
      "    )\n",
      "    (3): MBConv(\n",
      "      (block): Sequential(\n",
      "        (0): Conv2dNormActivation(\n",
      "          (0): Conv2d(192, 1152, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(1152, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (1): Conv2dNormActivation(\n",
      "          (0): Conv2d(1152, 1152, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), groups=1152, bias=False)\n",
      "          (1): BatchNorm2d(1152, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (2): SqueezeExcitation(\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
      "          (fc1): Conv2d(1152, 48, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (fc2): Conv2d(48, 1152, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (activation): SiLU(inplace=True)\n",
      "          (scale_activation): Sigmoid()\n",
      "        )\n",
      "        (3): Conv2dNormActivation(\n",
      "          (0): Conv2d(1152, 192, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (stochastic_depth): StochasticDepth(p=0.17500000000000002, mode=row)\n",
      "    )\n",
      "  )\n",
      "  (7): Sequential(\n",
      "    (0): MBConv(\n",
      "      (block): Sequential(\n",
      "        (0): Conv2dNormActivation(\n",
      "          (0): Conv2d(192, 1152, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(1152, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (1): Conv2dNormActivation(\n",
      "          (0): Conv2d(1152, 1152, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1152, bias=False)\n",
      "          (1): BatchNorm2d(1152, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): SiLU(inplace=True)\n",
      "        )\n",
      "        (2): SqueezeExcitation(\n",
      "          (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
      "          (fc1): Conv2d(1152, 48, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (fc2): Conv2d(48, 1152, kernel_size=(1, 1), stride=(1, 1))\n",
      "          (activation): SiLU(inplace=True)\n",
      "          (scale_activation): Sigmoid()\n",
      "        )\n",
      "        (3): Conv2dNormActivation(\n",
      "          (0): Conv2d(1152, 320, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(320, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (stochastic_depth): StochasticDepth(p=0.1875, mode=row)\n",
      "    )\n",
      "  )\n",
      "  (8): Conv2dNormActivation(\n",
      "    (0): Conv2d(320, 1280, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "    (1): BatchNorm2d(1280, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (2): SiLU(inplace=True)\n",
      "  )\n",
      ")\n",
      "*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-\n",
      "1 AdaptiveAvgPool2d(output_size=1)\n",
      "*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-\n",
      "2 Sequential(\n",
      "  (0): Dropout(p=0.2, inplace=True)\n",
      "  (1): Linear(in_features=1280, out_features=1000, bias=True)\n",
      ")\n",
      "*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-\n"
     ]
    }
   ],
   "source": [
    "for i, layer in enumerate(base_model.children()):\n",
    "    print(i , layer)\n",
    "    print(\"*-\"*20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "6dac0c40-33d3-4278-ad95-3faf198c81f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_layers = nn.Sequential()\n",
    "for i, layer in enumerate(base_model.children()):\n",
    "    if i <2:\n",
    "        base_layers.add_module(str(i), layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50dbacd7-b6a6-4f1a-bf2b-f2ce4c16b649",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "05d4166d-2d7d-4c79-a87c-9366953343ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Baseline_Model(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.base_model_headless = base_layers\n",
    "        self.base_model_headless.requires_grad = False\n",
    "        \n",
    "        self.fc1 = nn.Linear(1280, 512)\n",
    "        self.fc2 = nn.Linear(512, 128)\n",
    "        self.fc3 = nn.Linear(128, 5)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.dropout = nn.Dropout(p = 0.1)\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.base_model_headless(x)\n",
    "        # print(\"Output shape after Passing through Pretrained Model\")\n",
    "        # print(x.shape)\n",
    "\n",
    "        # Flatten\n",
    "        x = x.view((x.shape[0], -1))\n",
    "\n",
    "        # FCNN\n",
    "        x = self.fc1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.dropout(x)\n",
    "        x = self.fc2(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.dropout(x)\n",
    "        x = self.fc3(x)\n",
    "        # print(x.shape)\n",
    "\n",
    "        return x      "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "1935fc38-4ab5-47b9-94c6-091a1c9cdf5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/pytorch/vision/issues/7744\n",
    "\n",
    "from torchvision.models._api import WeightsEnum\n",
    "from torch.hub import load_state_dict_from_url\n",
    "\n",
    "def get_state_dict(self, *args, **kwargs):\n",
    "    kwargs.pop(\"check_hash\")\n",
    "    return load_state_dict_from_url(self.url, *args, **kwargs)\n",
    "WeightsEnum.get_state_dict = get_state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a65a8e15-615a-4c7d-b339-3882a0148b0c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "1a15c053-9780-47d1-882f-15fbbb06ac7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize Model\n",
    "model = Baseline_Model().to(CONFIG[\"device\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "6f956533-8f0b-4dc3-bbc6-03ac4975d84d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [-1, 32, 112, 112]             864\n",
      "       BatchNorm2d-2         [-1, 32, 112, 112]              64\n",
      "              SiLU-3         [-1, 32, 112, 112]               0\n",
      "            Conv2d-4         [-1, 32, 112, 112]             288\n",
      "       BatchNorm2d-5         [-1, 32, 112, 112]              64\n",
      "              SiLU-6         [-1, 32, 112, 112]               0\n",
      " AdaptiveAvgPool2d-7             [-1, 32, 1, 1]               0\n",
      "            Conv2d-8              [-1, 8, 1, 1]             264\n",
      "              SiLU-9              [-1, 8, 1, 1]               0\n",
      "           Conv2d-10             [-1, 32, 1, 1]             288\n",
      "          Sigmoid-11             [-1, 32, 1, 1]               0\n",
      "SqueezeExcitation-12         [-1, 32, 112, 112]               0\n",
      "           Conv2d-13         [-1, 16, 112, 112]             512\n",
      "      BatchNorm2d-14         [-1, 16, 112, 112]              32\n",
      "           MBConv-15         [-1, 16, 112, 112]               0\n",
      "           Conv2d-16         [-1, 96, 112, 112]           1,536\n",
      "      BatchNorm2d-17         [-1, 96, 112, 112]             192\n",
      "             SiLU-18         [-1, 96, 112, 112]               0\n",
      "           Conv2d-19           [-1, 96, 56, 56]             864\n",
      "      BatchNorm2d-20           [-1, 96, 56, 56]             192\n",
      "             SiLU-21           [-1, 96, 56, 56]               0\n",
      "AdaptiveAvgPool2d-22             [-1, 96, 1, 1]               0\n",
      "           Conv2d-23              [-1, 4, 1, 1]             388\n",
      "             SiLU-24              [-1, 4, 1, 1]               0\n",
      "           Conv2d-25             [-1, 96, 1, 1]             480\n",
      "          Sigmoid-26             [-1, 96, 1, 1]               0\n",
      "SqueezeExcitation-27           [-1, 96, 56, 56]               0\n",
      "           Conv2d-28           [-1, 24, 56, 56]           2,304\n",
      "      BatchNorm2d-29           [-1, 24, 56, 56]              48\n",
      "           MBConv-30           [-1, 24, 56, 56]               0\n",
      "           Conv2d-31          [-1, 144, 56, 56]           3,456\n",
      "      BatchNorm2d-32          [-1, 144, 56, 56]             288\n",
      "             SiLU-33          [-1, 144, 56, 56]               0\n",
      "           Conv2d-34          [-1, 144, 56, 56]           1,296\n",
      "      BatchNorm2d-35          [-1, 144, 56, 56]             288\n",
      "             SiLU-36          [-1, 144, 56, 56]               0\n",
      "AdaptiveAvgPool2d-37            [-1, 144, 1, 1]               0\n",
      "           Conv2d-38              [-1, 6, 1, 1]             870\n",
      "             SiLU-39              [-1, 6, 1, 1]               0\n",
      "           Conv2d-40            [-1, 144, 1, 1]           1,008\n",
      "          Sigmoid-41            [-1, 144, 1, 1]               0\n",
      "SqueezeExcitation-42          [-1, 144, 56, 56]               0\n",
      "           Conv2d-43           [-1, 24, 56, 56]           3,456\n",
      "      BatchNorm2d-44           [-1, 24, 56, 56]              48\n",
      "  StochasticDepth-45           [-1, 24, 56, 56]               0\n",
      "           MBConv-46           [-1, 24, 56, 56]               0\n",
      "           Conv2d-47          [-1, 144, 56, 56]           3,456\n",
      "      BatchNorm2d-48          [-1, 144, 56, 56]             288\n",
      "             SiLU-49          [-1, 144, 56, 56]               0\n",
      "           Conv2d-50          [-1, 144, 28, 28]           3,600\n",
      "      BatchNorm2d-51          [-1, 144, 28, 28]             288\n",
      "             SiLU-52          [-1, 144, 28, 28]               0\n",
      "AdaptiveAvgPool2d-53            [-1, 144, 1, 1]               0\n",
      "           Conv2d-54              [-1, 6, 1, 1]             870\n",
      "             SiLU-55              [-1, 6, 1, 1]               0\n",
      "           Conv2d-56            [-1, 144, 1, 1]           1,008\n",
      "          Sigmoid-57            [-1, 144, 1, 1]               0\n",
      "SqueezeExcitation-58          [-1, 144, 28, 28]               0\n",
      "           Conv2d-59           [-1, 40, 28, 28]           5,760\n",
      "      BatchNorm2d-60           [-1, 40, 28, 28]              80\n",
      "           MBConv-61           [-1, 40, 28, 28]               0\n",
      "           Conv2d-62          [-1, 240, 28, 28]           9,600\n",
      "      BatchNorm2d-63          [-1, 240, 28, 28]             480\n",
      "             SiLU-64          [-1, 240, 28, 28]               0\n",
      "           Conv2d-65          [-1, 240, 28, 28]           6,000\n",
      "      BatchNorm2d-66          [-1, 240, 28, 28]             480\n",
      "             SiLU-67          [-1, 240, 28, 28]               0\n",
      "AdaptiveAvgPool2d-68            [-1, 240, 1, 1]               0\n",
      "           Conv2d-69             [-1, 10, 1, 1]           2,410\n",
      "             SiLU-70             [-1, 10, 1, 1]               0\n",
      "           Conv2d-71            [-1, 240, 1, 1]           2,640\n",
      "          Sigmoid-72            [-1, 240, 1, 1]               0\n",
      "SqueezeExcitation-73          [-1, 240, 28, 28]               0\n",
      "           Conv2d-74           [-1, 40, 28, 28]           9,600\n",
      "      BatchNorm2d-75           [-1, 40, 28, 28]              80\n",
      "  StochasticDepth-76           [-1, 40, 28, 28]               0\n",
      "           MBConv-77           [-1, 40, 28, 28]               0\n",
      "           Conv2d-78          [-1, 240, 28, 28]           9,600\n",
      "      BatchNorm2d-79          [-1, 240, 28, 28]             480\n",
      "             SiLU-80          [-1, 240, 28, 28]               0\n",
      "           Conv2d-81          [-1, 240, 14, 14]           2,160\n",
      "      BatchNorm2d-82          [-1, 240, 14, 14]             480\n",
      "             SiLU-83          [-1, 240, 14, 14]               0\n",
      "AdaptiveAvgPool2d-84            [-1, 240, 1, 1]               0\n",
      "           Conv2d-85             [-1, 10, 1, 1]           2,410\n",
      "             SiLU-86             [-1, 10, 1, 1]               0\n",
      "           Conv2d-87            [-1, 240, 1, 1]           2,640\n",
      "          Sigmoid-88            [-1, 240, 1, 1]               0\n",
      "SqueezeExcitation-89          [-1, 240, 14, 14]               0\n",
      "           Conv2d-90           [-1, 80, 14, 14]          19,200\n",
      "      BatchNorm2d-91           [-1, 80, 14, 14]             160\n",
      "           MBConv-92           [-1, 80, 14, 14]               0\n",
      "           Conv2d-93          [-1, 480, 14, 14]          38,400\n",
      "      BatchNorm2d-94          [-1, 480, 14, 14]             960\n",
      "             SiLU-95          [-1, 480, 14, 14]               0\n",
      "           Conv2d-96          [-1, 480, 14, 14]           4,320\n",
      "      BatchNorm2d-97          [-1, 480, 14, 14]             960\n",
      "             SiLU-98          [-1, 480, 14, 14]               0\n",
      "AdaptiveAvgPool2d-99            [-1, 480, 1, 1]               0\n",
      "          Conv2d-100             [-1, 20, 1, 1]           9,620\n",
      "            SiLU-101             [-1, 20, 1, 1]               0\n",
      "          Conv2d-102            [-1, 480, 1, 1]          10,080\n",
      "         Sigmoid-103            [-1, 480, 1, 1]               0\n",
      "SqueezeExcitation-104          [-1, 480, 14, 14]               0\n",
      "          Conv2d-105           [-1, 80, 14, 14]          38,400\n",
      "     BatchNorm2d-106           [-1, 80, 14, 14]             160\n",
      " StochasticDepth-107           [-1, 80, 14, 14]               0\n",
      "          MBConv-108           [-1, 80, 14, 14]               0\n",
      "          Conv2d-109          [-1, 480, 14, 14]          38,400\n",
      "     BatchNorm2d-110          [-1, 480, 14, 14]             960\n",
      "            SiLU-111          [-1, 480, 14, 14]               0\n",
      "          Conv2d-112          [-1, 480, 14, 14]           4,320\n",
      "     BatchNorm2d-113          [-1, 480, 14, 14]             960\n",
      "            SiLU-114          [-1, 480, 14, 14]               0\n",
      "AdaptiveAvgPool2d-115            [-1, 480, 1, 1]               0\n",
      "          Conv2d-116             [-1, 20, 1, 1]           9,620\n",
      "            SiLU-117             [-1, 20, 1, 1]               0\n",
      "          Conv2d-118            [-1, 480, 1, 1]          10,080\n",
      "         Sigmoid-119            [-1, 480, 1, 1]               0\n",
      "SqueezeExcitation-120          [-1, 480, 14, 14]               0\n",
      "          Conv2d-121           [-1, 80, 14, 14]          38,400\n",
      "     BatchNorm2d-122           [-1, 80, 14, 14]             160\n",
      " StochasticDepth-123           [-1, 80, 14, 14]               0\n",
      "          MBConv-124           [-1, 80, 14, 14]               0\n",
      "          Conv2d-125          [-1, 480, 14, 14]          38,400\n",
      "     BatchNorm2d-126          [-1, 480, 14, 14]             960\n",
      "            SiLU-127          [-1, 480, 14, 14]               0\n",
      "          Conv2d-128          [-1, 480, 14, 14]          12,000\n",
      "     BatchNorm2d-129          [-1, 480, 14, 14]             960\n",
      "            SiLU-130          [-1, 480, 14, 14]               0\n",
      "AdaptiveAvgPool2d-131            [-1, 480, 1, 1]               0\n",
      "          Conv2d-132             [-1, 20, 1, 1]           9,620\n",
      "            SiLU-133             [-1, 20, 1, 1]               0\n",
      "          Conv2d-134            [-1, 480, 1, 1]          10,080\n",
      "         Sigmoid-135            [-1, 480, 1, 1]               0\n",
      "SqueezeExcitation-136          [-1, 480, 14, 14]               0\n",
      "          Conv2d-137          [-1, 112, 14, 14]          53,760\n",
      "     BatchNorm2d-138          [-1, 112, 14, 14]             224\n",
      "          MBConv-139          [-1, 112, 14, 14]               0\n",
      "          Conv2d-140          [-1, 672, 14, 14]          75,264\n",
      "     BatchNorm2d-141          [-1, 672, 14, 14]           1,344\n",
      "            SiLU-142          [-1, 672, 14, 14]               0\n",
      "          Conv2d-143          [-1, 672, 14, 14]          16,800\n",
      "     BatchNorm2d-144          [-1, 672, 14, 14]           1,344\n",
      "            SiLU-145          [-1, 672, 14, 14]               0\n",
      "AdaptiveAvgPool2d-146            [-1, 672, 1, 1]               0\n",
      "          Conv2d-147             [-1, 28, 1, 1]          18,844\n",
      "            SiLU-148             [-1, 28, 1, 1]               0\n",
      "          Conv2d-149            [-1, 672, 1, 1]          19,488\n",
      "         Sigmoid-150            [-1, 672, 1, 1]               0\n",
      "SqueezeExcitation-151          [-1, 672, 14, 14]               0\n",
      "          Conv2d-152          [-1, 112, 14, 14]          75,264\n",
      "     BatchNorm2d-153          [-1, 112, 14, 14]             224\n",
      " StochasticDepth-154          [-1, 112, 14, 14]               0\n",
      "          MBConv-155          [-1, 112, 14, 14]               0\n",
      "          Conv2d-156          [-1, 672, 14, 14]          75,264\n",
      "     BatchNorm2d-157          [-1, 672, 14, 14]           1,344\n",
      "            SiLU-158          [-1, 672, 14, 14]               0\n",
      "          Conv2d-159          [-1, 672, 14, 14]          16,800\n",
      "     BatchNorm2d-160          [-1, 672, 14, 14]           1,344\n",
      "            SiLU-161          [-1, 672, 14, 14]               0\n",
      "AdaptiveAvgPool2d-162            [-1, 672, 1, 1]               0\n",
      "          Conv2d-163             [-1, 28, 1, 1]          18,844\n",
      "            SiLU-164             [-1, 28, 1, 1]               0\n",
      "          Conv2d-165            [-1, 672, 1, 1]          19,488\n",
      "         Sigmoid-166            [-1, 672, 1, 1]               0\n",
      "SqueezeExcitation-167          [-1, 672, 14, 14]               0\n",
      "          Conv2d-168          [-1, 112, 14, 14]          75,264\n",
      "     BatchNorm2d-169          [-1, 112, 14, 14]             224\n",
      " StochasticDepth-170          [-1, 112, 14, 14]               0\n",
      "          MBConv-171          [-1, 112, 14, 14]               0\n",
      "          Conv2d-172          [-1, 672, 14, 14]          75,264\n",
      "     BatchNorm2d-173          [-1, 672, 14, 14]           1,344\n",
      "            SiLU-174          [-1, 672, 14, 14]               0\n",
      "          Conv2d-175            [-1, 672, 7, 7]          16,800\n",
      "     BatchNorm2d-176            [-1, 672, 7, 7]           1,344\n",
      "            SiLU-177            [-1, 672, 7, 7]               0\n",
      "AdaptiveAvgPool2d-178            [-1, 672, 1, 1]               0\n",
      "          Conv2d-179             [-1, 28, 1, 1]          18,844\n",
      "            SiLU-180             [-1, 28, 1, 1]               0\n",
      "          Conv2d-181            [-1, 672, 1, 1]          19,488\n",
      "         Sigmoid-182            [-1, 672, 1, 1]               0\n",
      "SqueezeExcitation-183            [-1, 672, 7, 7]               0\n",
      "          Conv2d-184            [-1, 192, 7, 7]         129,024\n",
      "     BatchNorm2d-185            [-1, 192, 7, 7]             384\n",
      "          MBConv-186            [-1, 192, 7, 7]               0\n",
      "          Conv2d-187           [-1, 1152, 7, 7]         221,184\n",
      "     BatchNorm2d-188           [-1, 1152, 7, 7]           2,304\n",
      "            SiLU-189           [-1, 1152, 7, 7]               0\n",
      "          Conv2d-190           [-1, 1152, 7, 7]          28,800\n",
      "     BatchNorm2d-191           [-1, 1152, 7, 7]           2,304\n",
      "            SiLU-192           [-1, 1152, 7, 7]               0\n",
      "AdaptiveAvgPool2d-193           [-1, 1152, 1, 1]               0\n",
      "          Conv2d-194             [-1, 48, 1, 1]          55,344\n",
      "            SiLU-195             [-1, 48, 1, 1]               0\n",
      "          Conv2d-196           [-1, 1152, 1, 1]          56,448\n",
      "         Sigmoid-197           [-1, 1152, 1, 1]               0\n",
      "SqueezeExcitation-198           [-1, 1152, 7, 7]               0\n",
      "          Conv2d-199            [-1, 192, 7, 7]         221,184\n",
      "     BatchNorm2d-200            [-1, 192, 7, 7]             384\n",
      " StochasticDepth-201            [-1, 192, 7, 7]               0\n",
      "          MBConv-202            [-1, 192, 7, 7]               0\n",
      "          Conv2d-203           [-1, 1152, 7, 7]         221,184\n",
      "     BatchNorm2d-204           [-1, 1152, 7, 7]           2,304\n",
      "            SiLU-205           [-1, 1152, 7, 7]               0\n",
      "          Conv2d-206           [-1, 1152, 7, 7]          28,800\n",
      "     BatchNorm2d-207           [-1, 1152, 7, 7]           2,304\n",
      "            SiLU-208           [-1, 1152, 7, 7]               0\n",
      "AdaptiveAvgPool2d-209           [-1, 1152, 1, 1]               0\n",
      "          Conv2d-210             [-1, 48, 1, 1]          55,344\n",
      "            SiLU-211             [-1, 48, 1, 1]               0\n",
      "          Conv2d-212           [-1, 1152, 1, 1]          56,448\n",
      "         Sigmoid-213           [-1, 1152, 1, 1]               0\n",
      "SqueezeExcitation-214           [-1, 1152, 7, 7]               0\n",
      "          Conv2d-215            [-1, 192, 7, 7]         221,184\n",
      "     BatchNorm2d-216            [-1, 192, 7, 7]             384\n",
      " StochasticDepth-217            [-1, 192, 7, 7]               0\n",
      "          MBConv-218            [-1, 192, 7, 7]               0\n",
      "          Conv2d-219           [-1, 1152, 7, 7]         221,184\n",
      "     BatchNorm2d-220           [-1, 1152, 7, 7]           2,304\n",
      "            SiLU-221           [-1, 1152, 7, 7]               0\n",
      "          Conv2d-222           [-1, 1152, 7, 7]          28,800\n",
      "     BatchNorm2d-223           [-1, 1152, 7, 7]           2,304\n",
      "            SiLU-224           [-1, 1152, 7, 7]               0\n",
      "AdaptiveAvgPool2d-225           [-1, 1152, 1, 1]               0\n",
      "          Conv2d-226             [-1, 48, 1, 1]          55,344\n",
      "            SiLU-227             [-1, 48, 1, 1]               0\n",
      "          Conv2d-228           [-1, 1152, 1, 1]          56,448\n",
      "         Sigmoid-229           [-1, 1152, 1, 1]               0\n",
      "SqueezeExcitation-230           [-1, 1152, 7, 7]               0\n",
      "          Conv2d-231            [-1, 192, 7, 7]         221,184\n",
      "     BatchNorm2d-232            [-1, 192, 7, 7]             384\n",
      " StochasticDepth-233            [-1, 192, 7, 7]               0\n",
      "          MBConv-234            [-1, 192, 7, 7]               0\n",
      "          Conv2d-235           [-1, 1152, 7, 7]         221,184\n",
      "     BatchNorm2d-236           [-1, 1152, 7, 7]           2,304\n",
      "            SiLU-237           [-1, 1152, 7, 7]               0\n",
      "          Conv2d-238           [-1, 1152, 7, 7]          10,368\n",
      "     BatchNorm2d-239           [-1, 1152, 7, 7]           2,304\n",
      "            SiLU-240           [-1, 1152, 7, 7]               0\n",
      "AdaptiveAvgPool2d-241           [-1, 1152, 1, 1]               0\n",
      "          Conv2d-242             [-1, 48, 1, 1]          55,344\n",
      "            SiLU-243             [-1, 48, 1, 1]               0\n",
      "          Conv2d-244           [-1, 1152, 1, 1]          56,448\n",
      "         Sigmoid-245           [-1, 1152, 1, 1]               0\n",
      "SqueezeExcitation-246           [-1, 1152, 7, 7]               0\n",
      "          Conv2d-247            [-1, 320, 7, 7]         368,640\n",
      "     BatchNorm2d-248            [-1, 320, 7, 7]             640\n",
      "          MBConv-249            [-1, 320, 7, 7]               0\n",
      "          Conv2d-250           [-1, 1280, 7, 7]         409,600\n",
      "     BatchNorm2d-251           [-1, 1280, 7, 7]           2,560\n",
      "            SiLU-252           [-1, 1280, 7, 7]               0\n",
      "AdaptiveAvgPool2d-253           [-1, 1280, 1, 1]               0\n",
      "          Linear-254                  [-1, 512]         655,872\n",
      "            ReLU-255                  [-1, 512]               0\n",
      "         Dropout-256                  [-1, 512]               0\n",
      "          Linear-257                  [-1, 128]          65,664\n",
      "            ReLU-258                  [-1, 128]               0\n",
      "         Dropout-259                  [-1, 128]               0\n",
      "          Linear-260                    [-1, 5]             645\n",
      "================================================================\n",
      "Total params: 4,729,729\n",
      "Trainable params: 4,729,729\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.57\n",
      "Forward/backward pass size (MB): 173.65\n",
      "Params size (MB): 18.04\n",
      "Estimated Total Size (MB): 192.27\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "summary(model, input_size = (3,224,224))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "310f1bca-1e3a-47fe-b02c-7a324e3a2aa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainloader = DataLoader(train_dataset, batch_size = CONFIG['batch_size'], shuffle = True, drop_last = True, num_workers = 2)\n",
    "\n",
    "testloader = DataLoader(test_dataset, batch_size = CONFIG['batch_size'], shuffle = False, drop_last = True, num_workers = 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc581391-85ef-4c5f-a278-df23c638a460",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "572a3216-d0f7-453f-9d18-209661e60591",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "44a71bce-186f-4c0c-9b32-308f9819afc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_one_epoch(model, trainloader, device, optimizer, criterion):\n",
    "    \"\"\"\n",
    "      Train 1 epoch on trainloader.\"\n",
    "    \"\"\"\n",
    "\n",
    "    model = model.train()\n",
    "    \n",
    "    # Keep track of loss and other evaluation metrics\n",
    "    train_loss = 0.0\n",
    "\n",
    "    YPredict = []\n",
    "    YTrue = []\n",
    "\n",
    "    ## Loop over all the batches\n",
    "    for i, (images, labels) in tqdm(enumerate(trainloader, 1), total=len(trainloader), desc=f\"Training on whole batch...\"):\n",
    "            \n",
    "        # Move images and labels to `device` (CPU or GPU)\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "\n",
    "        logits = model(images)\n",
    "\n",
    "        # print(logits)\n",
    "        # print(labels)\n",
    "        # print(logits.dtype, labels.dtype)\n",
    "        # print(type(labels[0]))\n",
    "        # print(type(labels))\n",
    "        loss = criterion(logits, labels)\n",
    "        # loss = criterion(logits, labels.to(torch.long))\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # Compute loss and accuracy for this batch\n",
    "        train_loss += loss.detach().item()\n",
    "\n",
    "        predictions = torch.argmax(logits, dim = 1)\n",
    "        predictions = predictions.view(-1)\n",
    "        YTrue.extend(labels.tolist())\n",
    "        YPredict.extend(predictions.tolist())\n",
    "\n",
    "    # Compute Scores\n",
    "    accuracy = accuracy_score(YTrue, YPredict)\n",
    "    f1 = f1_score(YTrue, YPredict, average = \"macro\")\n",
    "\n",
    "    return (train_loss/i, accuracy, f1) ## avg loss and acc over all batches\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb41d874-840c-46ba-a732-b538cff27f27",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "fd3de3db-b023-41f9-b2d9-cf9f8c245a37",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use `torch.no_grad()` here to disable gradient calculation. \n",
    "# It will reduce memory consumption as we don't need to compute gradients in inference.\n",
    "\n",
    "@torch.no_grad()\n",
    "def evaluate(model, testloader, criterion, device):\n",
    "    \n",
    "    model.eval()\n",
    "    \n",
    "    # Keep track of loss and other evaluation metrics\n",
    "    test_loss = 0.0\n",
    "\n",
    "    YPredict = []\n",
    "    YTrue = []\n",
    "    \n",
    "    # Loop through each batch on test set\n",
    "    for i, (images, labels) in enumerate(testloader, 1):\n",
    "        \n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "        \n",
    "        logits = model(images)\n",
    "\n",
    "        loss = criterion(logits, labels)\n",
    "        \n",
    "        test_loss += loss.detach().item()\n",
    "\n",
    "        predictions = torch.argmax(logits, dim = 1)\n",
    "        predictions = predictions.view(-1)\n",
    "        YTrue.extend(labels.tolist())\n",
    "        YPredict.extend(predictions.tolist())\n",
    "\n",
    "    # Compute Scores\n",
    "    accuracy = accuracy_score(YTrue, YPredict)\n",
    "    f1 = f1_score(YTrue, YPredict, average = \"macro\")\n",
    "\n",
    "    return (test_loss/i, accuracy, f1) ## avg loss and acc over all batches\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "9857e5e7-7f67-45c6-90f9-9a6586b650c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, trainloader, testloader, CONFIG):\n",
    "    \n",
    "    optimizer = AdamW(model.parameters(), lr = 0.01, weight_decay = 1e-5)\n",
    "\n",
    "    # criterion = nn.CrossEntropyLoss(weight = class_weights_tensor)\n",
    "\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "\n",
    "    ## Measure runtime\n",
    "    t_start = time()\n",
    "    \n",
    "    ## Store training log\n",
    "    history = defaultdict(list)\n",
    "    \n",
    "    # We will train the model `num_epochs` times\n",
    "    for i in range(1, CONFIG[\"epochs\"]+1):\n",
    "\n",
    "        # Training\n",
    "        train_loss, train_acc, train_f1 = train_one_epoch(model = model, trainloader = trainloader, optimizer = optimizer, criterion = criterion, device = CONFIG[\"device\"])\n",
    "\n",
    "        # Model Checkpointing\n",
    "        torch.save(model, f'models/Baseline_model/epoch{i}_batch_size{CONFIG[\"batch_size\"]}.pth')\n",
    "\n",
    "        # Evaluate on test set\n",
    "        test_loss, test_acc, test_f1 = evaluate(model = model, testloader = testloader, criterion = criterion, device = CONFIG[\"device\"])\n",
    "\n",
    "        ## Store train/test loss, accuracy, f1 scores\n",
    "        history[\"train_loss\"].append(train_loss)\n",
    "        history[\"train_acc\"].append(train_acc)\n",
    "        history[\"train_f1\"].append(train_f1)\n",
    "        \n",
    "        history[\"test_loss\"].append(test_loss)\n",
    "        history[\"test_acc\"].append(test_acc)\n",
    "        history[\"test_f1\"].append(test_f1)\n",
    "\n",
    "        \n",
    "        \n",
    "        ## print out train/test loss, accuracy\n",
    "        print(f'Epoch: {i} | Runtime: {((time()-t_start)/60):.2f}[m] | train_loss: {train_loss:.3f} | test_loss: {test_loss:.3f} | train_acc: {train_acc:.3f} |  test_acc: {test_acc:.3f} | train_f1: {train_f1:.3f} | test_f1: {test_f1:.3f}')    \n",
    "    return history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84b0e68d-1d4a-4885-92ee-29a9abf92b05",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training on whole batch...: 100%|██████████| 37/37 [07:07<00:00, 11.54s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1 | Runtime: 10.27[m] | train_loss: 1.829 | test_loss: 8.623 | train_acc: 0.314 |  test_acc: 0.336 | train_f1: 0.191 | test_f1: 0.126\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training on whole batch...: 100%|██████████| 37/37 [07:30<00:00, 12.18s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 2 | Runtime: 20.97[m] | train_loss: 1.528 | test_loss: 1.545 | train_acc: 0.321 |  test_acc: 0.203 | train_f1: 0.133 | test_f1: 0.068\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training on whole batch...: 100%|██████████| 37/37 [07:04<00:00, 11.49s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 3 | Runtime: 31.24[m] | train_loss: 1.558 | test_loss: 1.454 | train_acc: 0.338 |  test_acc: 0.484 | train_f1: 0.162 | test_f1: 0.131\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training on whole batch...: 100%|██████████| 37/37 [07:39<00:00, 12.43s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 4 | Runtime: 42.09[m] | train_loss: 1.537 | test_loss: 1.467 | train_acc: 0.382 |  test_acc: 0.484 | train_f1: 0.127 | test_f1: 0.131\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training on whole batch...: 100%|██████████| 37/37 [07:10<00:00, 11.64s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 5 | Runtime: 52.44[m] | train_loss: 1.508 | test_loss: 1.495 | train_acc: 0.378 |  test_acc: 0.438 | train_f1: 0.110 | test_f1: 0.141\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training on whole batch...: 100%|██████████| 37/37 [07:14<00:00, 11.76s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 6 | Runtime: 62.86[m] | train_loss: 1.504 | test_loss: 1.440 | train_acc: 0.372 |  test_acc: 0.484 | train_f1: 0.139 | test_f1: 0.131\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training on whole batch...: 100%|██████████| 37/37 [07:07<00:00, 11.56s/it]\n"
     ]
    }
   ],
   "source": [
    "# Train Model\n",
    "\n",
    "\n",
    "history_model = train_model(model, trainloader, testloader, CONFIG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad0b4e03-74a6-4f50-ab92-66ae5087fb36",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "fe441e3b-ca49-4169-9237-25c32dd4891b",
   "metadata": {},
   "source": [
    "# Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d6bb5c3-1575-43e9-94b0-35e272d47f45",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Plot train_loss, test_loss\n",
    "num_epochs = CONFIG['epochs']\n",
    "plt.plot(np.arange(num_epochs), history_model[\"train_loss\"], label='train loss')\n",
    "plt.plot(np.arange(num_epochs), history_model[\"test_loss\"], label='test loss')\n",
    "plt.xlabel('#epochs')\n",
    "plt.ylabel('loss')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7060a727-2c6c-4528-ba8b-ca08f4c4123e",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Plot accuracy\n",
    "\n",
    "plt.plot(np.arange(num_epochs), history_model[\"train_acc\"], label='train acc')\n",
    "plt.plot(np.arange(num_epochs), history_model[\"test_acc\"], label='test acc')\n",
    "plt.xlabel('#epochs')\n",
    "plt.ylabel('accuracy')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82512b88-1775-4423-832e-7f1d865de437",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Plot accuracy\n",
    "\n",
    "plt.plot(np.arange(num_epochs), history_model[\"train_f1\"], label='train f1')\n",
    "plt.plot(np.arange(num_epochs), history_model[\"test_f1\"], label='test f1')\n",
    "plt.xlabel('#epochs')\n",
    "plt.ylabel('f1')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71b5f7d9-95af-4002-9b4e-f0a520d09dcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch.save(model.state_dict(), f'models/Baseline_model/state_dict_epoch{i}_batch_size{CONFIG[\"batch_size\"]}.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d66a740e-4beb-46e0-8208-cf3d4da05782",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}