pytorch_mobilenet_(1).ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "RXw5og44LNSk",
    "outputId": "0986447e-feae-40e3-8755-198227496909"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mounted at /content/drive\n"
     ]
    }
   ],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive',force_remount=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n",
      "cuda:0\n"
     ]
    }
   ],
   "source": [
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)\n",
    "x = torch.tensor([1.0, 2.0, 3.0])\n",
    "x = x.to(device)\n",
    "print(x.device) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.device_count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting torch\n",
      "  Using cached torch-2.0.1-cp39-cp39-manylinux1_x86_64.whl (619.9 MB)\n",
      "Collecting torchvision\n",
      "  Using cached torchvision-0.15.2-cp39-cp39-manylinux1_x86_64.whl (6.0 MB)\n",
      "Collecting filelock (from torch)\n",
      "  Obtaining dependency information for filelock from https://files.pythonhosted.org/packages/00/45/ec3407adf6f6b5bf867a4462b2b0af27597a26bd3cd6e2534cb6ab029938/filelock-3.12.2-py3-none-any.whl.metadata\n",
      "  Downloading filelock-3.12.2-py3-none-any.whl.metadata (2.7 kB)\n",
      "Requirement already satisfied: typing-extensions in ./anaconda3/envs/mobilenet/lib/python3.9/site-packages (from torch) (4.5.0)\n",
      "Collecting sympy (from torch)\n",
      "  Downloading sympy-1.12-py3-none-any.whl (5.7 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.7/5.7 MB\u001b[0m \u001b[31m42.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hCollecting networkx (from torch)\n",
      "  Downloading networkx-3.1-py3-none-any.whl (2.1 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.1/2.1 MB\u001b[0m \u001b[31m83.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: jinja2 in ./anaconda3/envs/mobilenet/lib/python3.9/site-packages (from torch) (3.1.2)\n",
      "Collecting nvidia-cuda-nvrtc-cu11==11.7.99 (from torch)\n",
      "  Downloading nvidia_cuda_nvrtc_cu11-11.7.99-2-py3-none-manylinux1_x86_64.whl (21.0 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m21.0/21.0 MB\u001b[0m \u001b[31m65.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-cuda-runtime-cu11==11.7.99 (from torch)\n",
      "  Downloading nvidia_cuda_runtime_cu11-11.7.99-py3-none-manylinux1_x86_64.whl (849 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m849.3/849.3 kB\u001b[0m \u001b[31m60.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-cuda-cupti-cu11==11.7.101 (from torch)\n",
      "  Downloading nvidia_cuda_cupti_cu11-11.7.101-py3-none-manylinux1_x86_64.whl (11.8 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m11.8/11.8 MB\u001b[0m \u001b[31m74.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-cudnn-cu11==8.5.0.96 (from torch)\n",
      "  Downloading nvidia_cudnn_cu11-8.5.0.96-2-py3-none-manylinux1_x86_64.whl (557.1 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m557.1/557.1 MB\u001b[0m \u001b[31m7.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-cublas-cu11==11.10.3.66 (from torch)\n",
      "  Downloading nvidia_cublas_cu11-11.10.3.66-py3-none-manylinux1_x86_64.whl (317.1 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m317.1/317.1 MB\u001b[0m \u001b[31m11.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-cufft-cu11==10.9.0.58 (from torch)\n",
      "  Downloading nvidia_cufft_cu11-10.9.0.58-py3-none-manylinux1_x86_64.whl (168.4 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m168.4/168.4 MB\u001b[0m \u001b[31m19.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-curand-cu11==10.2.10.91 (from torch)\n",
      "  Downloading nvidia_curand_cu11-10.2.10.91-py3-none-manylinux1_x86_64.whl (54.6 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m54.6/54.6 MB\u001b[0m \u001b[31m43.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-cusolver-cu11==11.4.0.1 (from torch)\n",
      "  Downloading nvidia_cusolver_cu11-11.4.0.1-2-py3-none-manylinux1_x86_64.whl (102.6 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m102.6/102.6 MB\u001b[0m \u001b[31m27.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-cusparse-cu11==11.7.4.91 (from torch)\n",
      "  Downloading nvidia_cusparse_cu11-11.7.4.91-py3-none-manylinux1_x86_64.whl (173.2 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m173.2/173.2 MB\u001b[0m \u001b[31m19.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-nccl-cu11==2.14.3 (from torch)\n",
      "  Downloading nvidia_nccl_cu11-2.14.3-py3-none-manylinux1_x86_64.whl (177.1 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m177.1/177.1 MB\u001b[0m \u001b[31m18.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting nvidia-nvtx-cu11==11.7.91 (from torch)\n",
      "  Downloading nvidia_nvtx_cu11-11.7.91-py3-none-manylinux1_x86_64.whl (98 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m98.6/98.6 kB\u001b[0m \u001b[31m15.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting triton==2.0.0 (from torch)\n",
      "  Downloading triton-2.0.0-1-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (63.3 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m63.3/63.3 MB\u001b[0m \u001b[31m14.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: setuptools in ./anaconda3/envs/mobilenet/lib/python3.9/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch) (67.8.0)\n",
      "Requirement already satisfied: wheel in ./anaconda3/envs/mobilenet/lib/python3.9/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch) (0.38.4)\n",
      "Collecting cmake (from triton==2.0.0->torch)\n",
      "  Obtaining dependency information for cmake from https://files.pythonhosted.org/packages/14/b8/06f8fdc4687af3d3d8d95461d97737df2f144acd28eff65a3c47c29d0152/cmake-3.27.0-py2.py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata\n",
      "  Downloading cmake-3.27.0-py2.py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (6.7 kB)\n",
      "Collecting lit (from triton==2.0.0->torch)\n",
      "  Downloading lit-16.0.6.tar.gz (153 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m153.7/153.7 kB\u001b[0m \u001b[31m22.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Installing build dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n",
      "\u001b[?25h  Installing backend dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Preparing metadata (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied: numpy in ./anaconda3/envs/mobilenet/lib/python3.9/site-packages (from torchvision) (1.24.3)\n",
      "Requirement already satisfied: requests in ./anaconda3/envs/mobilenet/lib/python3.9/site-packages (from torchvision) (2.31.0)\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in ./anaconda3/envs/mobilenet/lib/python3.9/site-packages (from torchvision) (10.0.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in ./anaconda3/envs/mobilenet/lib/python3.9/site-packages (from jinja2->torch) (2.1.3)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in ./anaconda3/envs/mobilenet/lib/python3.9/site-packages (from requests->torchvision) (3.2.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in ./anaconda3/envs/mobilenet/lib/python3.9/site-packages (from requests->torchvision) (3.4)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in ./anaconda3/envs/mobilenet/lib/python3.9/site-packages (from requests->torchvision) (1.26.16)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in ./anaconda3/envs/mobilenet/lib/python3.9/site-packages (from requests->torchvision) (2023.5.7)\n",
      "Collecting mpmath>=0.19 (from sympy->torch)\n",
      "  Downloading mpmath-1.3.0-py3-none-any.whl (536 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m536.2/536.2 kB\u001b[0m \u001b[31m35.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading filelock-3.12.2-py3-none-any.whl (10 kB)\n",
      "Downloading cmake-3.27.0-py2.py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (26.0 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m26.0/26.0 MB\u001b[0m \u001b[31m60.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m:00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hBuilding wheels for collected packages: lit\n",
      "  Building wheel for lit (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for lit: filename=lit-16.0.6-py3-none-any.whl size=93582 sha256=d3b7e5b0b63b7a1fb043393f8ea9c7ce9efa9ed2ca2fbd27ed4a9c3c6fba2fa5\n",
      "  Stored in directory: /home/kailashv/.cache/pip/wheels/a5/36/d6/cac2e6fb891889b33a548f2fddb8b4b7726399aaa2ed32b188\n",
      "Successfully built lit\n",
      "Installing collected packages: mpmath, lit, cmake, sympy, nvidia-nvtx-cu11, nvidia-nccl-cu11, nvidia-cusparse-cu11, nvidia-curand-cu11, nvidia-cufft-cu11, nvidia-cuda-runtime-cu11, nvidia-cuda-nvrtc-cu11, nvidia-cuda-cupti-cu11, nvidia-cublas-cu11, networkx, filelock, nvidia-cusolver-cu11, nvidia-cudnn-cu11, triton, torch, torchvision\n",
      "  Attempting uninstall: nvidia-cublas-cu11\n",
      "    Found existing installation: nvidia-cublas-cu11 11.11.3.6\n",
      "    Uninstalling nvidia-cublas-cu11-11.11.3.6:\n",
      "      Successfully uninstalled nvidia-cublas-cu11-11.11.3.6\n",
      "  Attempting uninstall: nvidia-cudnn-cu11\n",
      "    Found existing installation: nvidia-cudnn-cu11 8.6.0.163\n",
      "    Uninstalling nvidia-cudnn-cu11-8.6.0.163:\n",
      "      Successfully uninstalled nvidia-cudnn-cu11-8.6.0.163\n",
      "Successfully installed cmake-3.27.0 filelock-3.12.2 lit-16.0.6 mpmath-1.3.0 networkx-3.1 nvidia-cublas-cu11-11.10.3.66 nvidia-cuda-cupti-cu11-11.7.101 nvidia-cuda-nvrtc-cu11-11.7.99 nvidia-cuda-runtime-cu11-11.7.99 nvidia-cudnn-cu11-8.5.0.96 nvidia-cufft-cu11-10.9.0.58 nvidia-curand-cu11-10.2.10.91 nvidia-cusolver-cu11-11.4.0.1 nvidia-cusparse-cu11-11.7.4.91 nvidia-nccl-cu11-2.14.3 nvidia-nvtx-cu11-11.7.91 sympy-1.12 torch-2.0.1 torchvision-0.15.2 triton-2.0.0\n"
     ]
    }
   ],
   "source": [
    "!pip install torch torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "id": "npmDe_CGKoPO"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.models as models\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpu_id = 0  # Change this to the index of the desired GPU\n",
    "torch.cuda.set_device(gpu_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "import zipfile\n",
    "zip_path = 'Updated Data.zip'\n",
    "with zipfile.ZipFile(zip_path,'r') as zip_ref:\n",
    "    zip_ref.extractall('RediMinds/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "id": "rGAyI9xdKxTB"
   },
   "outputs": [],
   "source": [
    "DATADIR = \"RediMinds/Data_Cleaned/\"\n",
    "categories = [\"Anaesthesia_machine\",\"baby_incubator\",\"Bone_saws\",\"C_arm\",\"colonoscope\",\"Curved_Mayo_Scissor\",\"difibrillator\",\"Electrocautery_devices\",\"fetal_doppler\",\"forceps\",\"Heart_Lung_Machine\",\"IABP\",\"IMRT\",\"infusion_pump\",\"Laryngoscopes\",\"mayfield_clamp\",\"Needle_Biopsy_Device\",\"phacoemulsification\",\"Radiofrequency_Ablation_Device\",\"Scalpel\",\"Straight_Dissection_Clamp\",\"Straight_Mayo_Scissor\",\"Suction_Machine\",\"ventilator\",\"x_ray\"]\n",
    "input_size = (224, 224)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "25"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(categories)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "editable": true,
    "id": "d0mpzTyLK1l2",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, data_dir, categories, input_size, transform=None):\n",
    "        self.data_dir = data_dir\n",
    "        self.categories = categories\n",
    "        self.input_size = input_size\n",
    "        self.transform = transform\n",
    "        self.data = []\n",
    "        self.labels = []\n",
    "        self.load_data()\n",
    "\n",
    "    def load_data(self):\n",
    "        for cat in self.categories:\n",
    "            path = os.path.join(self.data_dir, cat)\n",
    "            class_num = self.categories.index(cat)\n",
    "            for img_name in os.listdir(path):\n",
    "                try:\n",
    "                    img_path = os.path.join(path, img_name)\n",
    "                    img = Image.open(img_path).convert(\"RGB\")\n",
    "                    img = img.resize(self.input_size)\n",
    "                    if self.transform:\n",
    "                        img = self.transform(img)\n",
    "                    self.data.append(img)\n",
    "                    self.labels.append(class_num)\n",
    "                except Exception as e:\n",
    "                    pass\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img = self.data[idx]\n",
    "        label = self.labels[idx]\n",
    "        return img, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4y-YObECLeC2",
    "outputId": "37531a35-eaa5-45a9-a0d6-b7a5e3d683ad"
   },
   "outputs": [],
   "source": [
    "transform = transforms.Compose([\n",
    "    transforms.RandomResizedCrop(input_size, scale=(0.8, 1.0)),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    #transforms.RandomRotation(15),\n",
    "    #transforms.ColorJitter(brightness=0.2, contrast = 0.25, saturation = 0.2, hue=0.1),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.2, 0.2, 0.2))\n",
    "])\n",
    "\n",
    "dataset = CustomDataset(DATADIR, categories, input_size, transform=transform)\n",
    "\n",
    "train_size = int(0.70 * len(dataset))\n",
    "validation_size = int(0.15 * len(dataset))\n",
    "test_size = len(dataset) - train_size - validation_size\n",
    "train_dataset, validation_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size,validation_size, test_size])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "id": "HDjJkmtgLjVZ"
   },
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "validation_loader = DataLoader(validation_dataset, batch_size=batch_size)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Q-YSuecIN8Nd"
   },
   "source": [
    "class CustomMobileNet(nn.Module):\n",
    "    def __init__(self, num_classes, fine_tune=True):\n",
    "        super(CustomMobileNet, self).__init__()\n",
    "        self.base_model = models.mobilenet_v3_large(pretrained='imagenet')\n",
    "        if not fine_tune:\n",
    "            for param in self.base_model.parameters():\n",
    "                param.requires_grad = False\n",
    "\n",
    "        num_features = self.base_model.classifier[0].in_features\n",
    "        self.base_model.classifier = nn.Sequential(\n",
    "            nn.Conv2d(num_features, 128, kernel_size=1),\n",
    "            nn.AdaptiveAvgPool2d(1))\n",
    "        self.cnn1 = nn.Conv2d(960,2048,kernel_size=3, padding=1)\n",
    "        self.dropout1 = nn.Dropout(0.5)\n",
    "        #self.cnn2 = nn.Conv2d(1024,512,kernel_size = 3, padding=1)\n",
    "        #self.dropout5 = nn.Dropout(0.3)\n",
    "        self.fc1 = nn.Linear(2048, 2048)\n",
    "        self.dropout2 = nn.Dropout(0.5)\n",
    "        self.fc7 = nn.Linear(2048, 1024)\n",
    "        self.dropout5 = nn.Dropout(0.5)\n",
    "        self.fc8 = nn.Linear(1024,1024)\n",
    "        self.dropout6 = nn.Dropout(0.4)\n",
    "        self.fc6 = nn.Linear(1024, num_classes)\n",
    "        # self.fc2 = nn.Linear(128, 64)\n",
    "        # self.dropout3 = nn.Dropout(0.4)\n",
    "        # self.fc3 = nn.Linear(64, 32)\n",
    "        # self.dropout4 = nn.Dropout(0.2)\n",
    "        # self.fc4 = nn.Linear(32, 16)\n",
    "        # self.fc5 = nn.Linear(16, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.base_model.features(x)\n",
    "        x = self.cnn1(x)\n",
    "        x = torch.relu(x)\n",
    "        x = self.dropout1(x)\n",
    "        #x = self.cnn2(x)\n",
    "        #x = torch.relu(x)\n",
    "        x = torch.mean(x, dim=[2, 3])  # Global average pooling\n",
    "        x = x.view(x.size(0), -1)\n",
    "        #x = self.dropout5(x)\n",
    "        x = self.fc1(x)\n",
    "        x = torch.relu(x)\n",
    "        x = self.dropout2(x)\n",
    "        x = self.fc7(x)\n",
    "        x = torch.relu(x)\n",
    "        x = self.dropout5(x)\n",
    "        x = self.fc8(x)\n",
    "        x = torch.relu(x)\n",
    "        x = self.dropout6(x)\n",
    "        x = self.fc6(x)\n",
    "        # x = torch.relu(x)\n",
    "        # x = self.fc2(x)\n",
    "        # x = torch.relu(x)\n",
    "        # x = self.dropout3(x)\n",
    "        # x = self.fc3(x)\n",
    "        # x = torch.relu(x)\n",
    "        # # x = self.dropout4(x)\n",
    "        # x = self.fc4(x)\n",
    "        # x = torch.relu(x)\n",
    "        # x = self.fc5(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomMobileNet(nn.Module):\n",
    "    def __init__(self, num_classes,fine_tune = False):\n",
    "        super(CustomMobileNet, self).__init__()\n",
    "        self.base_model = models.mobilenet_v3_large(pretrained=models.MobileNet_V3_Large_Weights.IMAGENET1K_V2)\n",
    "        self.base_model = models.mobilenet_v3_large(pretrained=models.MobileNet_V3_Large_Weights.IMAGENET1K_V2)\n",
    "        \n",
    "\n",
    "        if not fine_tune:\n",
    "            # Freeze all layers except the last few\n",
    "            for param in self.base_model.parameters():\n",
    "                param.requires_grad = False\n",
    "\n",
    "            # Identify the last few layers of the base model and set requires_grad to True\n",
    "            for param in self.base_model.features[-3:].parameters():\n",
    "                param.requires_grad = True\n",
    "                \n",
    "        num_features = self.base_model.classifier[-1].in_features\n",
    "        self.base_model.classifier[-1] = nn.Sequential(\n",
    "            nn.Conv2d(num_features, 128, kernel_size=1),\n",
    "            nn.ReLU(),\n",
    "            nn.AdaptiveAvgPool2d(1))\n",
    "        self.cnn1 = nn.Conv2d(960,4098,kernel_size=3, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(4098)\n",
    "        self.dropout1 = nn.Dropout(0.6)\n",
    "        # self.cnn2 = nn.Conv2d(1024, 1024,kernel_size = 3, padding = 1)\n",
    "        # self.dropout2 = nn.Dropout(0.6)\n",
    "        # self.cnn3 = nn.Conv2d(1024, 512,kernel_size = 3, padding = 1)\n",
    "        # self.dropout3 = nn.Dropout(0.5)\n",
    "        # self.bn2 = nn.BatchNorm2d(512)\n",
    "        self.cnn4 = nn.Conv2d(4098, 4098,kernel_size = 3, padding = 1)\n",
    "        self.bn4 = nn.BatchNorm2d(4098)\n",
    "        self.dropout4 = nn.Dropout(0.5)\n",
    "        self.cnn5 = nn.Conv2d(4098, 2048,kernel_size = 3, padding = 1)\n",
    "        self.dropout5 = nn.Dropout(0.6)\n",
    "        self.bn3 = nn.BatchNorm2d(2048)\n",
    "        self.avgpooling = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        self.base_model = nn.Sequential(*list(self.base_model.children())[:-1])\n",
    "        # self.fc1 = nn.Linear(128, 64)\n",
    "        # self.dropout6 = nn.Dropout(0.4)\n",
    "        self.fc2 = nn.Linear(2048,2048)\n",
    "        self.bnl = nn.BatchNorm1d(2048)\n",
    "        self.dropout7 = nn.Dropout(0.5)\n",
    "        self.fc3 = nn.Linear(2048, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.base_model(x)\n",
    "        x = self.cnn1(x)\n",
    "        x = self.bn1(x)\n",
    "        x = torch.relu(x)\n",
    "        x = self.dropout1(x)\n",
    "        # x = self.cnn2(x)\n",
    "        # x = torch.relu(x)\n",
    "        # x = self.dropout2(x)\n",
    "        # x = self.cnn3(x)\n",
    "        # x = self.bn2(x)\n",
    "        # x = torch.relu(x)\n",
    "        # x = self.dropout3(x)\n",
    "        # x = self.cnn4(x)\n",
    "        # x = self.bn4(x)\n",
    "        # x = torch.relu(x)\n",
    "        # x = self.dropout4(x)\n",
    "        x = self.cnn5(x)\n",
    "        x = self.bn3(x)\n",
    "        x = torch.relu(x)\n",
    "        x = self.dropout5(x)\n",
    "        x = self.avgpooling(x)\n",
    "        #x = torch.mean(x, dim=[2, 3])  # Global average pooling\n",
    "        x = x.view(x.size(0), -1)\n",
    "        # x = self.fc1(x)\n",
    "        # x = torch.relu(x)\n",
    "        # x = self.dropout6(x)\n",
    "        x = self.fc2(x)\n",
    "        x = self.bnl(x)\n",
    "        x = torch.relu(x)\n",
    "        x = self.dropout7(x)\n",
    "        x = self.fc3(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "weight_enum = torch.hub.load(\"pytorch/vision\", \"get_model_weights\", name=\"mobilenet_v3_large\")\n",
    "print([weight for weight in weight_enum])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {
    "editable": true,
    "id": "x03pIT-qMgF1",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from torch.optim.lr_scheduler import StepLR\n",
    "num_classes = len(categories)\n",
    "model = CustomMobileNet(num_classes).to(device)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.005,weight_decay=1e-6)\n",
    "scheduler = StepLR(optimizer, step_size = 20, gamma=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "base_model.0.0.0.weight: requires_grad=False\n",
      "base_model.0.0.1.weight: requires_grad=False\n",
      "base_model.0.0.1.bias: requires_grad=False\n",
      "base_model.0.1.block.0.0.weight: requires_grad=False\n",
      "base_model.0.1.block.0.1.weight: requires_grad=False\n",
      "base_model.0.1.block.0.1.bias: requires_grad=False\n",
      "base_model.0.1.block.1.0.weight: requires_grad=False\n",
      "base_model.0.1.block.1.1.weight: requires_grad=False\n",
      "base_model.0.1.block.1.1.bias: requires_grad=False\n",
      "base_model.0.2.block.0.0.weight: requires_grad=False\n",
      "base_model.0.2.block.0.1.weight: requires_grad=False\n",
      "base_model.0.2.block.0.1.bias: requires_grad=False\n",
      "base_model.0.2.block.1.0.weight: requires_grad=False\n",
      "base_model.0.2.block.1.1.weight: requires_grad=False\n",
      "base_model.0.2.block.1.1.bias: requires_grad=False\n",
      "base_model.0.2.block.2.0.weight: requires_grad=False\n",
      "base_model.0.2.block.2.1.weight: requires_grad=False\n",
      "base_model.0.2.block.2.1.bias: requires_grad=False\n",
      "base_model.0.3.block.0.0.weight: requires_grad=False\n",
      "base_model.0.3.block.0.1.weight: requires_grad=False\n",
      "base_model.0.3.block.0.1.bias: requires_grad=False\n",
      "base_model.0.3.block.1.0.weight: requires_grad=False\n",
      "base_model.0.3.block.1.1.weight: requires_grad=False\n",
      "base_model.0.3.block.1.1.bias: requires_grad=False\n",
      "base_model.0.3.block.2.0.weight: requires_grad=False\n",
      "base_model.0.3.block.2.1.weight: requires_grad=False\n",
      "base_model.0.3.block.2.1.bias: requires_grad=False\n",
      "base_model.0.4.block.0.0.weight: requires_grad=False\n",
      "base_model.0.4.block.0.1.weight: requires_grad=False\n",
      "base_model.0.4.block.0.1.bias: requires_grad=False\n",
      "base_model.0.4.block.1.0.weight: requires_grad=False\n",
      "base_model.0.4.block.1.1.weight: requires_grad=False\n",
      "base_model.0.4.block.1.1.bias: requires_grad=False\n",
      "base_model.0.4.block.2.fc1.weight: requires_grad=False\n",
      "base_model.0.4.block.2.fc1.bias: requires_grad=False\n",
      "base_model.0.4.block.2.fc2.weight: requires_grad=False\n",
      "base_model.0.4.block.2.fc2.bias: requires_grad=False\n",
      "base_model.0.4.block.3.0.weight: requires_grad=False\n",
      "base_model.0.4.block.3.1.weight: requires_grad=False\n",
      "base_model.0.4.block.3.1.bias: requires_grad=False\n",
      "base_model.0.5.block.0.0.weight: requires_grad=False\n",
      "base_model.0.5.block.0.1.weight: requires_grad=False\n",
      "base_model.0.5.block.0.1.bias: requires_grad=False\n",
      "base_model.0.5.block.1.0.weight: requires_grad=False\n",
      "base_model.0.5.block.1.1.weight: requires_grad=False\n",
      "base_model.0.5.block.1.1.bias: requires_grad=False\n",
      "base_model.0.5.block.2.fc1.weight: requires_grad=False\n",
      "base_model.0.5.block.2.fc1.bias: requires_grad=False\n",
      "base_model.0.5.block.2.fc2.weight: requires_grad=False\n",
      "base_model.0.5.block.2.fc2.bias: requires_grad=False\n",
      "base_model.0.5.block.3.0.weight: requires_grad=False\n",
      "base_model.0.5.block.3.1.weight: requires_grad=False\n",
      "base_model.0.5.block.3.1.bias: requires_grad=False\n",
      "base_model.0.6.block.0.0.weight: requires_grad=False\n",
      "base_model.0.6.block.0.1.weight: requires_grad=False\n",
      "base_model.0.6.block.0.1.bias: requires_grad=False\n",
      "base_model.0.6.block.1.0.weight: requires_grad=False\n",
      "base_model.0.6.block.1.1.weight: requires_grad=False\n",
      "base_model.0.6.block.1.1.bias: requires_grad=False\n",
      "base_model.0.6.block.2.fc1.weight: requires_grad=False\n",
      "base_model.0.6.block.2.fc1.bias: requires_grad=False\n",
      "base_model.0.6.block.2.fc2.weight: requires_grad=False\n",
      "base_model.0.6.block.2.fc2.bias: requires_grad=False\n",
      "base_model.0.6.block.3.0.weight: requires_grad=False\n",
      "base_model.0.6.block.3.1.weight: requires_grad=False\n",
      "base_model.0.6.block.3.1.bias: requires_grad=False\n",
      "base_model.0.7.block.0.0.weight: requires_grad=False\n",
      "base_model.0.7.block.0.1.weight: requires_grad=False\n",
      "base_model.0.7.block.0.1.bias: requires_grad=False\n",
      "base_model.0.7.block.1.0.weight: requires_grad=False\n",
      "base_model.0.7.block.1.1.weight: requires_grad=False\n",
      "base_model.0.7.block.1.1.bias: requires_grad=False\n",
      "base_model.0.7.block.2.0.weight: requires_grad=False\n",
      "base_model.0.7.block.2.1.weight: requires_grad=False\n",
      "base_model.0.7.block.2.1.bias: requires_grad=False\n",
      "base_model.0.8.block.0.0.weight: requires_grad=False\n",
      "base_model.0.8.block.0.1.weight: requires_grad=False\n",
      "base_model.0.8.block.0.1.bias: requires_grad=False\n",
      "base_model.0.8.block.1.0.weight: requires_grad=False\n",
      "base_model.0.8.block.1.1.weight: requires_grad=False\n",
      "base_model.0.8.block.1.1.bias: requires_grad=False\n",
      "base_model.0.8.block.2.0.weight: requires_grad=False\n",
      "base_model.0.8.block.2.1.weight: requires_grad=False\n",
      "base_model.0.8.block.2.1.bias: requires_grad=False\n",
      "base_model.0.9.block.0.0.weight: requires_grad=False\n",
      "base_model.0.9.block.0.1.weight: requires_grad=False\n",
      "base_model.0.9.block.0.1.bias: requires_grad=False\n",
      "base_model.0.9.block.1.0.weight: requires_grad=False\n",
      "base_model.0.9.block.1.1.weight: requires_grad=False\n",
      "base_model.0.9.block.1.1.bias: requires_grad=False\n",
      "base_model.0.9.block.2.0.weight: requires_grad=False\n",
      "base_model.0.9.block.2.1.weight: requires_grad=False\n",
      "base_model.0.9.block.2.1.bias: requires_grad=False\n",
      "base_model.0.10.block.0.0.weight: requires_grad=False\n",
      "base_model.0.10.block.0.1.weight: requires_grad=False\n",
      "base_model.0.10.block.0.1.bias: requires_grad=False\n",
      "base_model.0.10.block.1.0.weight: requires_grad=False\n",
      "base_model.0.10.block.1.1.weight: requires_grad=False\n",
      "base_model.0.10.block.1.1.bias: requires_grad=False\n",
      "base_model.0.10.block.2.0.weight: requires_grad=False\n",
      "base_model.0.10.block.2.1.weight: requires_grad=False\n",
      "base_model.0.10.block.2.1.bias: requires_grad=False\n",
      "base_model.0.11.block.0.0.weight: requires_grad=False\n",
      "base_model.0.11.block.0.1.weight: requires_grad=False\n",
      "base_model.0.11.block.0.1.bias: requires_grad=False\n",
      "base_model.0.11.block.1.0.weight: requires_grad=False\n",
      "base_model.0.11.block.1.1.weight: requires_grad=False\n",
      "base_model.0.11.block.1.1.bias: requires_grad=False\n",
      "base_model.0.11.block.2.fc1.weight: requires_grad=False\n",
      "base_model.0.11.block.2.fc1.bias: requires_grad=False\n",
      "base_model.0.11.block.2.fc2.weight: requires_grad=False\n",
      "base_model.0.11.block.2.fc2.bias: requires_grad=False\n",
      "base_model.0.11.block.3.0.weight: requires_grad=False\n",
      "base_model.0.11.block.3.1.weight: requires_grad=False\n",
      "base_model.0.11.block.3.1.bias: requires_grad=False\n",
      "base_model.0.12.block.0.0.weight: requires_grad=False\n",
      "base_model.0.12.block.0.1.weight: requires_grad=False\n",
      "base_model.0.12.block.0.1.bias: requires_grad=False\n",
      "base_model.0.12.block.1.0.weight: requires_grad=False\n",
      "base_model.0.12.block.1.1.weight: requires_grad=False\n",
      "base_model.0.12.block.1.1.bias: requires_grad=False\n",
      "base_model.0.12.block.2.fc1.weight: requires_grad=False\n",
      "base_model.0.12.block.2.fc1.bias: requires_grad=False\n",
      "base_model.0.12.block.2.fc2.weight: requires_grad=False\n",
      "base_model.0.12.block.2.fc2.bias: requires_grad=False\n",
      "base_model.0.12.block.3.0.weight: requires_grad=False\n",
      "base_model.0.12.block.3.1.weight: requires_grad=False\n",
      "base_model.0.12.block.3.1.bias: requires_grad=False\n",
      "base_model.0.13.block.0.0.weight: requires_grad=False\n",
      "base_model.0.13.block.0.1.weight: requires_grad=False\n",
      "base_model.0.13.block.0.1.bias: requires_grad=False\n",
      "base_model.0.13.block.1.0.weight: requires_grad=False\n",
      "base_model.0.13.block.1.1.weight: requires_grad=False\n",
      "base_model.0.13.block.1.1.bias: requires_grad=False\n",
      "base_model.0.13.block.2.fc1.weight: requires_grad=False\n",
      "base_model.0.13.block.2.fc1.bias: requires_grad=False\n",
      "base_model.0.13.block.2.fc2.weight: requires_grad=False\n",
      "base_model.0.13.block.2.fc2.bias: requires_grad=False\n",
      "base_model.0.13.block.3.0.weight: requires_grad=False\n",
      "base_model.0.13.block.3.1.weight: requires_grad=False\n",
      "base_model.0.13.block.3.1.bias: requires_grad=False\n",
      "base_model.0.14.block.0.0.weight: requires_grad=True\n",
      "base_model.0.14.block.0.1.weight: requires_grad=True\n",
      "base_model.0.14.block.0.1.bias: requires_grad=True\n",
      "base_model.0.14.block.1.0.weight: requires_grad=True\n",
      "base_model.0.14.block.1.1.weight: requires_grad=True\n",
      "base_model.0.14.block.1.1.bias: requires_grad=True\n",
      "base_model.0.14.block.2.fc1.weight: requires_grad=True\n",
      "base_model.0.14.block.2.fc1.bias: requires_grad=True\n",
      "base_model.0.14.block.2.fc2.weight: requires_grad=True\n",
      "base_model.0.14.block.2.fc2.bias: requires_grad=True\n",
      "base_model.0.14.block.3.0.weight: requires_grad=True\n",
      "base_model.0.14.block.3.1.weight: requires_grad=True\n",
      "base_model.0.14.block.3.1.bias: requires_grad=True\n",
      "base_model.0.15.block.0.0.weight: requires_grad=True\n",
      "base_model.0.15.block.0.1.weight: requires_grad=True\n",
      "base_model.0.15.block.0.1.bias: requires_grad=True\n",
      "base_model.0.15.block.1.0.weight: requires_grad=True\n",
      "base_model.0.15.block.1.1.weight: requires_grad=True\n",
      "base_model.0.15.block.1.1.bias: requires_grad=True\n",
      "base_model.0.15.block.2.fc1.weight: requires_grad=True\n",
      "base_model.0.15.block.2.fc1.bias: requires_grad=True\n",
      "base_model.0.15.block.2.fc2.weight: requires_grad=True\n",
      "base_model.0.15.block.2.fc2.bias: requires_grad=True\n",
      "base_model.0.15.block.3.0.weight: requires_grad=True\n",
      "base_model.0.15.block.3.1.weight: requires_grad=True\n",
      "base_model.0.15.block.3.1.bias: requires_grad=True\n",
      "base_model.0.16.0.weight: requires_grad=True\n",
      "base_model.0.16.1.weight: requires_grad=True\n",
      "base_model.0.16.1.bias: requires_grad=True\n",
      "cnn1.weight: requires_grad=True\n",
      "cnn1.bias: requires_grad=True\n",
      "bn1.weight: requires_grad=True\n",
      "bn1.bias: requires_grad=True\n",
      "cnn4.weight: requires_grad=True\n",
      "cnn4.bias: requires_grad=True\n",
      "bn4.weight: requires_grad=True\n",
      "bn4.bias: requires_grad=True\n",
      "cnn5.weight: requires_grad=True\n",
      "cnn5.bias: requires_grad=True\n",
      "bn3.weight: requires_grad=True\n",
      "bn3.bias: requires_grad=True\n",
      "fc2.weight: requires_grad=True\n",
      "fc2.bias: requires_grad=True\n",
      "bnl.weight: requires_grad=True\n",
      "bnl.bias: requires_grad=True\n",
      "fc3.weight: requires_grad=True\n",
      "fc3.bias: requires_grad=True\n"
     ]
    }
   ],
   "source": [
    "for name, param in model.named_parameters():\n",
    "    print(f\"{name}: requires_grad={param.requires_grad}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0eNmu2UJMg7v",
    "outputId": "26814188-ecd3-4f26-bf3f-a7578768454f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/50 | Train Loss: 3.1270 | Train Acc: 0.2820 | Val Loss: 194.5615 | Val Acc: 0.0615\n",
      "Epoch 2/50 | Train Loss: 2.4702 | Train Acc: 0.4620 | Val Loss: 18.7272 | Val Acc: 0.1746\n",
      "Epoch 3/50 | Train Loss: 2.1446 | Train Acc: 0.5314 | Val Loss: 16.8683 | Val Acc: 0.3939\n",
      "Epoch 4/50 | Train Loss: 1.6966 | Train Acc: 0.5966 | Val Loss: 2.0972 | Val Acc: 0.5656\n",
      "Epoch 5/50 | Train Loss: 1.4978 | Train Acc: 0.6555 | Val Loss: 7.4148 | Val Acc: 0.4749\n",
      "Epoch 6/50 | Train Loss: 1.1915 | Train Acc: 0.7078 | Val Loss: 1.2464 | Val Acc: 0.6830\n",
      "Epoch 7/50 | Train Loss: 0.9551 | Train Acc: 0.7641 | Val Loss: 1.2442 | Val Acc: 0.7151\n",
      "Epoch 8/50 | Train Loss: 0.6976 | Train Acc: 0.7952 | Val Loss: 0.8147 | Val Acc: 0.7654\n"
     ]
    }
   ],
   "source": [
    "train_acc_history = []\n",
    "train_loss_history = []\n",
    "val_acc_history = []\n",
    "val_loss_history = []\n",
    "\n",
    "num_epochs = 50\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    train_loss, train_corrects = 0.0, 0.0\n",
    "\n",
    "    for images, labels in train_loader:\n",
    "        images, labels = images.to(device), labels.to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        train_loss += loss.item() * images.size(0)\n",
    "        _, preds = torch.max(outputs, 1)\n",
    "        train_corrects += torch.sum(preds == labels).item()\n",
    "\n",
    "    train_loss = train_loss / len(train_dataset)\n",
    "    train_acc = train_corrects / len(train_dataset)\n",
    "    train_acc_history.append(train_acc)\n",
    "    train_loss_history.append(train_loss)\n",
    "\n",
    "    model.eval()\n",
    "    val_loss, val_corrects = 0.0, 0.0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for images, labels in validation_loader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "\n",
    "            outputs = model(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "\n",
    "            val_loss += loss.item() * images.size(0)\n",
    "            _, preds = torch.max(outputs, 1)\n",
    "            val_corrects += torch.sum(preds == labels).item()\n",
    "\n",
    "    val_loss = val_loss / len(validation_dataset)\n",
    "    val_acc = val_corrects / len(validation_dataset)\n",
    "\n",
    "    val_acc_history.append(val_acc)\n",
    "    val_loss_history.append(val_loss)\n",
    "\n",
    "    scheduler.step()\n",
    "\n",
    "    print(f\"Epoch {epoch+1}/{num_epochs} | Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Plot training and validation accuracy\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.plot(range(1, num_epochs+1), train_acc_history, label='Train')\n",
    "plt.plot(range(1, num_epochs+1), val_acc_history, label='Validation')\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "# Plot training and validation loss\n",
    "plt.plot(range(1, num_epochs+1), train_loss_history, label='Train')\n",
    "plt.plot(range(1, num_epochs+1), val_loss_history, label='Validation')\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Loss')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "test_loss, test_corrects = 0.0, 0.0\n",
    "\n",
    "with torch.no_grad():\n",
    "    for images, labels in test_loader:\n",
    "        images, labels = images.to(device), labels.to(device)\n",
    "\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "\n",
    "        test_loss += loss.item() * images.size(0)\n",
    "        _, preds = torch.max(outputs, 1)\n",
    "        test_corrects += torch.sum(preds == labels).item()\n",
    "\n",
    "test_loss = test_loss / len(test_dataset)\n",
    "test_acc = test_corrects / len(test_dataset)\n",
    "\n",
    "print(f\"Test Loss: {test_loss:.4f} | Test Acc: {test_acc:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PFj0xxxMM4iY"
   },
   "source": [
    "model.eval()\n",
    "test_loss, test_corrects = 0.0, 0.0\n",
    "\n",
    "with torch.no_grad():\n",
    "    for images, labels in test_loader:\n",
    "        images, labels = images.to(device), labels.to(device)\n",
    "\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "\n",
    "        test_loss += loss.item() * images.size(0)\n",
    "        _, preds = torch.max(outputs, 1)\n",
    "        test_corrects += torch.sum(preds == labels).item()\n",
    "\n",
    "test_loss = test_loss / len(test_dataset)\n",
    "test_acc = test_corrects / len(test_dataset)\n",
    "\n",
    "print(f\"Test Loss: {test_loss:.4f} | Test Acc: {test_acc:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(),\"model_torch_mobilenet_2.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "def calculate_accuracy_and_loss(model, loader):\n",
    "    \"\"\"\n",
    "    Calculates the accuracy and loss for the given model and loader.\n",
    "\n",
    "    Args:\n",
    "        model: The model to evaluate.\n",
    "        loader: The data loader to use.\n",
    "\n",
    "    Returns:\n",
    "        The accuracy and loss values.\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    corrects = 0\n",
    "    total = 0\n",
    "    loss = 0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for images, labels in loader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "\n",
    "            outputs = model(images)\n",
    "            loss += criterion(outputs, labels).item() * images.size(0)\n",
    "\n",
    "            _, preds = torch.max(outputs, 1)\n",
    "            corrects += torch.sum(preds == labels).item()\n",
    "            total += labels.size(0)\n",
    "\n",
    "    return corrects / total, loss / total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "id": "2TY4aG8JSDgv",
    "jupyter": {
     "outputs_hidden": true,
     "source_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/50 | Train Loss: 3.4058 | Train Acc: 0.0711\n",
      "Test Loss: 3.2241 | Test Acc: 0.8341\n",
      "Epoch 2/50 | Train Loss: 3.2104 | Train Acc: 0.0495\n",
      "Test Loss: 3.2266 | Test Acc: 0.8341\n",
      "Epoch 3/50 | Train Loss: 3.2094 | Train Acc: 0.0509\n",
      "Test Loss: 3.2276 | Test Acc: 0.8341\n",
      "Epoch 4/50 | Train Loss: 3.1975 | Train Acc: 0.0513\n",
      "Test Loss: 3.2225 | Test Acc: 0.8341\n",
      "Epoch 5/50 | Train Loss: 3.1950 | Train Acc: 0.0505\n",
      "Test Loss: 3.2245 | Test Acc: 0.8341\n",
      "Epoch 6/50 | Train Loss: 3.1912 | Train Acc: 0.0493\n",
      "Test Loss: 3.2241 | Test Acc: 0.8341\n",
      "Epoch 7/50 | Train Loss: 3.1908 | Train Acc: 0.0509\n",
      "Test Loss: 3.2252 | Test Acc: 0.8341\n",
      "Epoch 8/50 | Train Loss: 3.1940 | Train Acc: 0.0515\n",
      "Test Loss: 3.2296 | Test Acc: 0.8341\n",
      "Epoch 9/50 | Train Loss: 3.1926 | Train Acc: 0.0515\n",
      "Test Loss: 3.2229 | Test Acc: 0.8341\n",
      "Epoch 10/50 | Train Loss: 3.1902 | Train Acc: 0.0499\n",
      "Test Loss: 3.2217 | Test Acc: 0.8341\n",
      "Epoch 11/50 | Train Loss: 3.1855 | Train Acc: 0.0519\n",
      "Test Loss: 3.2241 | Test Acc: 0.8341\n",
      "Epoch 12/50 | Train Loss: 3.1889 | Train Acc: 0.0499\n",
      "Test Loss: 3.2229 | Test Acc: 0.8341\n",
      "Epoch 13/50 | Train Loss: 3.1880 | Train Acc: 0.0521\n",
      "Test Loss: 3.2225 | Test Acc: 0.8341\n",
      "Epoch 14/50 | Train Loss: 3.1890 | Train Acc: 0.0501\n",
      "Test Loss: 3.2210 | Test Acc: 0.8341\n",
      "Epoch 15/50 | Train Loss: 3.1913 | Train Acc: 0.0487\n",
      "Test Loss: 3.2228 | Test Acc: 0.8341\n",
      "Epoch 16/50 | Train Loss: 3.1929 | Train Acc: 0.0545\n",
      "Test Loss: 3.2239 | Test Acc: 0.8341\n",
      "Epoch 17/50 | Train Loss: 3.1873 | Train Acc: 0.0547\n",
      "Test Loss: 3.2233 | Test Acc: 0.8341\n",
      "Epoch 18/50 | Train Loss: 3.1946 | Train Acc: 0.0499\n",
      "Test Loss: 3.2249 | Test Acc: 0.8341\n",
      "Epoch 19/50 | Train Loss: 3.1963 | Train Acc: 0.0495\n",
      "Test Loss: 3.2230 | Test Acc: 0.8341\n",
      "Epoch 20/50 | Train Loss: 3.1921 | Train Acc: 0.0564\n",
      "Test Loss: 3.2247 | Test Acc: 0.8341\n",
      "Epoch 21/50 | Train Loss: 3.2020 | Train Acc: 0.0541\n",
      "Test Loss: 3.2236 | Test Acc: 0.8341\n",
      "Epoch 22/50 | Train Loss: 3.2072 | Train Acc: 0.0483\n",
      "Test Loss: 3.2237 | Test Acc: 0.8341\n",
      "Epoch 23/50 | Train Loss: 3.2039 | Train Acc: 0.0473\n",
      "Test Loss: 3.2213 | Test Acc: 0.8341\n",
      "Epoch 24/50 | Train Loss: 3.1995 | Train Acc: 0.0477\n",
      "Test Loss: 3.2215 | Test Acc: 0.8341\n",
      "Epoch 25/50 | Train Loss: 3.1957 | Train Acc: 0.0541\n",
      "Test Loss: 3.2278 | Test Acc: 0.8341\n",
      "Epoch 26/50 | Train Loss: 3.2042 | Train Acc: 0.0509\n",
      "Test Loss: 3.2233 | Test Acc: 0.8341\n",
      "Epoch 27/50 | Train Loss: 3.2014 | Train Acc: 0.0448\n",
      "Test Loss: 3.2209 | Test Acc: 0.8341\n",
      "Epoch 28/50 | Train Loss: 3.2159 | Train Acc: 0.0499\n",
      "Test Loss: 3.2226 | Test Acc: 0.8341\n",
      "Epoch 29/50 | Train Loss: 3.1973 | Train Acc: 0.0487\n",
      "Test Loss: 3.2230 | Test Acc: 0.8341\n",
      "Epoch 30/50 | Train Loss: 3.2041 | Train Acc: 0.0479\n",
      "Test Loss: 3.2239 | Test Acc: 0.8341\n",
      "Epoch 31/50 | Train Loss: 3.2012 | Train Acc: 0.0465\n",
      "Test Loss: 3.2213 | Test Acc: 0.8341\n",
      "Epoch 32/50 | Train Loss: 3.1986 | Train Acc: 0.0491\n",
      "Test Loss: 3.2263 | Test Acc: 0.8341\n",
      "Epoch 33/50 | Train Loss: 3.1991 | Train Acc: 0.0521\n",
      "Test Loss: 3.2253 | Test Acc: 0.8341\n",
      "Epoch 34/50 | Train Loss: 3.1963 | Train Acc: 0.0491\n",
      "Test Loss: 3.2228 | Test Acc: 0.8341\n",
      "Epoch 35/50 | Train Loss: 3.1993 | Train Acc: 0.0475\n",
      "Test Loss: 3.2236 | Test Acc: 0.8341\n",
      "Epoch 36/50 | Train Loss: 3.2012 | Train Acc: 0.0485\n",
      "Test Loss: 3.2213 | Test Acc: 0.8341\n",
      "Epoch 37/50 | Train Loss: 3.1979 | Train Acc: 0.0471\n",
      "Test Loss: 3.2230 | Test Acc: 0.8341\n",
      "Epoch 38/50 | Train Loss: 3.1992 | Train Acc: 0.0493\n",
      "Test Loss: 3.2219 | Test Acc: 0.8341\n",
      "Epoch 39/50 | Train Loss: 3.1989 | Train Acc: 0.0491\n",
      "Test Loss: 3.2257 | Test Acc: 0.8341\n",
      "Epoch 40/50 | Train Loss: 3.1983 | Train Acc: 0.0483\n",
      "Test Loss: 3.2211 | Test Acc: 0.8341\n",
      "Epoch 41/50 | Train Loss: 3.2003 | Train Acc: 0.0495\n",
      "Test Loss: 3.2249 | Test Acc: 0.8341\n",
      "Epoch 42/50 | Train Loss: 3.1983 | Train Acc: 0.0467\n",
      "Test Loss: 3.2212 | Test Acc: 0.8341\n",
      "Epoch 43/50 | Train Loss: 3.2036 | Train Acc: 0.0458\n",
      "Test Loss: 3.2213 | Test Acc: 0.8341\n"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "train_accuracy_values = []\n",
    "test_accuracy_values = []\n",
    "train_loss_values = []\n",
    "test_loss_values = []\n",
    "\n",
    "num_epochs = 50\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    train_loss, train_corrects = 0.0, 0.0\n",
    "\n",
    "    for images, labels in train_loader:\n",
    "        images, labels = images.to(device), labels.to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        train_loss += loss.item() * images.size(0)\n",
    "        _, preds = torch.max(outputs, 1)\n",
    "        train_corrects += torch.sum(preds == labels).item()\n",
    "\n",
    "        train_accuracy_values.append(train_corrects / len(train_dataset))\n",
    "        train_loss_values.append(train_loss / len(train_dataset))\n",
    "\n",
    "        test_accuracy, test_loss = calculate_accuracy_and_loss(model, test_loader)\n",
    "        test_accuracy_values.append(test_accuracy)\n",
    "        test_loss_values.append(test_loss)\n",
    "\n",
    "    train_loss = train_loss / len(train_dataset)\n",
    "    train_acc = train_corrects / len(train_dataset)\n",
    "\n",
    "    print(f\"Epoch {epoch+1}/{num_epochs} | Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f}\")\n",
    "    print(f\"Test Loss: {test_loss:.4f} | Test Acc: {test_acc:.4f}\")\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(train_accuracy_values, label=\"Train Accuracy\")\n",
    "plt.plot(test_accuracy_values, label=\"Test Accuracy\")\n",
    "plt.legend()\n",
    "plt.title(\"Accuracy vs. Epoch\")\n",
    "plt.show()\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(train_loss_values, label=\"Train Loss\")\n",
    "plt.plot(test_loss_values, label=\"Test Loss\")\n",
    "plt.legend()\n",
    "plt.title(\"Loss vs. Epoch\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "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.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}