VGG16Model_pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "221f27ba-7d07-45d4-9e98-f59296b13b6c",
   "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": 4,
   "id": "44da099f-387e-426e-a70d-8e68ce85cc33",
   "metadata": {},
   "outputs": [],
   "source": [
    "gpu_id = 1  # Change this to the index of the desired GPU\n",
    "torch.cuda.set_device(gpu_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6db4bdbb-7706-44ed-af95-ac637f22f02a",
   "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": 6,
   "id": "8774e8bb-0cd0-486a-ba35-299d97df53da",
   "metadata": {
    "id": "d0mpzTyLK1l2"
   },
   "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": 7,
   "id": "a575d042-35ad-4178-827b-368da4f418e1",
   "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": 8,
   "id": "ed09a85b-0025-41cb-b87d-d42070b5fc42",
   "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",
   "id": "2bf5fcae-3af3-49f3-9dd6-88e8b286f0f4",
   "metadata": {},
   "source": [
    "class CustomVGG16(nn.Module):\n",
    "    def __init__(self, num_classes, fine_tune_layers=5):\n",
    "        super(CustomVGG16, self).__init__()\n",
    "        self.base_model = models.vgg16(pretrained=True)\n",
    "\n",
    "        for param in self.base_model.parameters():\n",
    "            param.requires_grad = False\n",
    "\n",
    "        num_features = self.base_model.classifier[6].in_features\n",
    "        self.base_model.classifier[6] = nn.Sequential(\n",
    "            nn.Conv2d(num_features, 128, kernel_size=1),\n",
    "            nn.AdaptiveAvgPool2d(1))\n",
    "\n",
    "        self.cnn1 = nn.Conv2d(512,1024,kernel_size=3, padding=1)\n",
    "        self.dropout1 = nn.Dropout(0.3)\n",
    "        self.cnn2 = nn.Conv2d(1024,256,kernel_size = 3, padding=1)\n",
    "        self.dropout2 = nn.Dropout(0.2)\n",
    "        self.cnn3 = nn.Conv2d(256,64, kernel_size = 1, padding = 1)\n",
    "        self.dropout3 = nn.Dropout(0.1)\n",
    "        self.cnnl = nn.Conv2d(64,num_classes, kernel_size=1)\n",
    "\n",
    "        layers_to_train = [\n",
    "            self.base_model.features[-fine_tune_layers:],\n",
    "            self.base_model.classifier[6],\n",
    "            self.cnn1,\n",
    "            self.cnn2,\n",
    "            self.cnn3,\n",
    "            self.cnnl]\n",
    "        \n",
    "        if fine_tune_layers > 0:\n",
    "            for layer in layers_to_train:\n",
    "                for param in layer.parameters():\n",
    "                    param.requires_grad = True\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 = self.cnn3(x)\n",
    "        x = torch.relu(x)\n",
    "        x = self.dropout3(x)\n",
    "        x = self.cnnl(x)\n",
    "        x = torch.flatten(x,1)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d0d36727-feda-4c02-b239-ea5504c34ec3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomVGG16(nn.Module):\n",
    "    def __init__(self, num_classes, fine_tune_layers=2):\n",
    "        super(CustomVGG16, self).__init__()\n",
    "        vgg16 = models.vgg16_bn(pretrained=models.VGG16_BN_Weights.IMAGENET1K_V1)\n",
    "\n",
    "        # Freeze all layers by default\n",
    "        for param in vgg16.parameters():\n",
    "            param.requires_grad = False\n",
    "\n",
    "        # Make the last block trainable\n",
    "        for layer in vgg16.features[-fine_tune_layers:]:\n",
    "            for param in layer.parameters():\n",
    "                param.requires_grad = True\n",
    "\n",
    "        for param in vgg16.classifier.parameters():\n",
    "            param.requires_grad = True\n",
    "\n",
    "        self.vgg16 = vgg16\n",
    "        \n",
    "        # Create the custom head (additional dense layers)\n",
    "        self.head = nn.Sequential(\n",
    "            nn.Linear(512 * 7 * 7, 4096),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.5),\n",
    "            nn.Linear(4096, 4096),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.5),\n",
    "            nn.Linear(4096,2048),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.5),\n",
    "            nn.Linear(2048, num_classes)\n",
    "        )\n",
    "        # self.cnn1 = nn.Conv2d(2048, 2048, kernel_size = 3, padding = 1)\n",
    "        # self.dropout = nn.Dropout(0.5)\n",
    "        # self.cnn2 = nn.Conv2d(2048, 512, kernel_size = 3, padding = 1)\n",
    "        # self.dropout1 = nn.Dropout(0.4)\n",
    "        # self.cnn3 = nn.Conv2d(512,64, kernel_size = 3, padding = 1)\n",
    "        # self.dropout2 = nn.Dropout(0.3)\n",
    "        # self.cnnl = nn.Conv2d(64, num_classes, kernel_size=3,padding=1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.vgg16.features(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = self.head(x)\n",
    "        # x = x.view(x.size(0), -1,1,1)\n",
    "        # x = self.cnn1(x)\n",
    "        # x = torch.relu(x)\n",
    "        # x = self.dropout(x)\n",
    "        # x = self.cnn2(x)\n",
    "        # x = torch.relu(x)\n",
    "        # x = self.dropout(x)\n",
    "        # x = self.cnn3(x)\n",
    "        # x = torch.relu(x)\n",
    "        # x = self.dropout2(x)\n",
    "        # x = self.cnnl(x)\n",
    "        x = torch.flatten(x,1)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "19f2898f-ebda-44fe-821c-aec6968371f5",
   "metadata": {
    "id": "x03pIT-qMgF1"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/kailashv/anaconda3/envs/mobilenet/lib/python3.9/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "/home/kailashv/anaconda3/envs/mobilenet/lib/python3.9/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=VGG16_BN_Weights.IMAGENET1K_V1`. You can also use `weights=VGG16_BN_Weights.DEFAULT` to get the most up-to-date weights.\n",
      "  warnings.warn(msg)\n"
     ]
    }
   ],
   "source": [
    "num_classes = len(categories)\n",
    "model = CustomVGG16(num_classes).to(device)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.0001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "7c0a5808-75ca-487f-8e5c-d5466d03b6aa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vgg16.features.0.weight: requires_grad=False\n",
      "vgg16.features.0.bias: requires_grad=False\n",
      "vgg16.features.1.weight: requires_grad=False\n",
      "vgg16.features.1.bias: requires_grad=False\n",
      "vgg16.features.3.weight: requires_grad=False\n",
      "vgg16.features.3.bias: requires_grad=False\n",
      "vgg16.features.4.weight: requires_grad=False\n",
      "vgg16.features.4.bias: requires_grad=False\n",
      "vgg16.features.7.weight: requires_grad=False\n",
      "vgg16.features.7.bias: requires_grad=False\n",
      "vgg16.features.8.weight: requires_grad=False\n",
      "vgg16.features.8.bias: requires_grad=False\n",
      "vgg16.features.10.weight: requires_grad=False\n",
      "vgg16.features.10.bias: requires_grad=False\n",
      "vgg16.features.11.weight: requires_grad=False\n",
      "vgg16.features.11.bias: requires_grad=False\n",
      "vgg16.features.14.weight: requires_grad=False\n",
      "vgg16.features.14.bias: requires_grad=False\n",
      "vgg16.features.15.weight: requires_grad=False\n",
      "vgg16.features.15.bias: requires_grad=False\n",
      "vgg16.features.17.weight: requires_grad=False\n",
      "vgg16.features.17.bias: requires_grad=False\n",
      "vgg16.features.18.weight: requires_grad=False\n",
      "vgg16.features.18.bias: requires_grad=False\n",
      "vgg16.features.20.weight: requires_grad=False\n",
      "vgg16.features.20.bias: requires_grad=False\n",
      "vgg16.features.21.weight: requires_grad=False\n",
      "vgg16.features.21.bias: requires_grad=False\n",
      "vgg16.features.24.weight: requires_grad=False\n",
      "vgg16.features.24.bias: requires_grad=False\n",
      "vgg16.features.25.weight: requires_grad=False\n",
      "vgg16.features.25.bias: requires_grad=False\n",
      "vgg16.features.27.weight: requires_grad=False\n",
      "vgg16.features.27.bias: requires_grad=False\n",
      "vgg16.features.28.weight: requires_grad=False\n",
      "vgg16.features.28.bias: requires_grad=False\n",
      "vgg16.features.30.weight: requires_grad=False\n",
      "vgg16.features.30.bias: requires_grad=False\n",
      "vgg16.features.31.weight: requires_grad=False\n",
      "vgg16.features.31.bias: requires_grad=False\n",
      "vgg16.features.34.weight: requires_grad=False\n",
      "vgg16.features.34.bias: requires_grad=False\n",
      "vgg16.features.35.weight: requires_grad=False\n",
      "vgg16.features.35.bias: requires_grad=False\n",
      "vgg16.features.37.weight: requires_grad=False\n",
      "vgg16.features.37.bias: requires_grad=False\n",
      "vgg16.features.38.weight: requires_grad=False\n",
      "vgg16.features.38.bias: requires_grad=False\n",
      "vgg16.features.40.weight: requires_grad=False\n",
      "vgg16.features.40.bias: requires_grad=False\n",
      "vgg16.features.41.weight: requires_grad=False\n",
      "vgg16.features.41.bias: requires_grad=False\n",
      "vgg16.classifier.0.weight: requires_grad=True\n",
      "vgg16.classifier.0.bias: requires_grad=True\n",
      "vgg16.classifier.3.weight: requires_grad=True\n",
      "vgg16.classifier.3.bias: requires_grad=True\n",
      "vgg16.classifier.6.weight: requires_grad=True\n",
      "vgg16.classifier.6.bias: requires_grad=True\n",
      "head.0.weight: requires_grad=True\n",
      "head.0.bias: requires_grad=True\n",
      "head.3.weight: requires_grad=True\n",
      "head.3.bias: requires_grad=True\n",
      "head.6.weight: requires_grad=True\n",
      "head.6.bias: requires_grad=True\n",
      "head.9.weight: requires_grad=True\n",
      "head.9.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": 19,
   "id": "d0c4dae9-bccc-48ed-a309-e2fd377755ca",
   "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: 2.8617 | Train Acc: 0.1465 | Val Loss: 2.1636 | Val Acc: 0.2891\n",
      "Epoch 2/50 | Train Loss: 1.7510 | Train Acc: 0.4417 | Val Loss: 1.3224 | Val Acc: 0.5908\n",
      "Epoch 3/50 | Train Loss: 0.9485 | Train Acc: 0.7022 | Val Loss: 1.0683 | Val Acc: 0.6564\n",
      "Epoch 4/50 | Train Loss: 0.4541 | Train Acc: 0.8559 | Val Loss: 0.9771 | Val Acc: 0.7137\n",
      "Epoch 5/50 | Train Loss: 0.2161 | Train Acc: 0.9360 | Val Loss: 0.9603 | Val Acc: 0.7263\n",
      "Epoch 6/50 | Train Loss: 0.0889 | Train Acc: 0.9758 | Val Loss: 1.0446 | Val Acc: 0.7304\n",
      "Epoch 7/50 | Train Loss: 0.0577 | Train Acc: 0.9836 | Val Loss: 1.0246 | Val Acc: 0.7374\n",
      "Epoch 8/50 | Train Loss: 0.0445 | Train Acc: 0.9874 | Val Loss: 1.1029 | Val Acc: 0.7263\n",
      "Epoch 9/50 | Train Loss: 0.0256 | Train Acc: 0.9955 | Val Loss: 1.0248 | Val Acc: 0.7570\n",
      "Epoch 10/50 | Train Loss: 0.0122 | Train Acc: 0.9988 | Val Loss: 1.0839 | Val Acc: 0.7570\n",
      "Epoch 11/50 | Train Loss: 0.0200 | Train Acc: 0.9949 | Val Loss: 1.1592 | Val Acc: 0.7318\n",
      "Epoch 12/50 | Train Loss: 0.0160 | Train Acc: 0.9961 | Val Loss: 1.2467 | Val Acc: 0.7318\n",
      "Epoch 13/50 | Train Loss: 0.0149 | Train Acc: 0.9955 | Val Loss: 1.3656 | Val Acc: 0.7137\n",
      "Epoch 14/50 | Train Loss: 0.0286 | Train Acc: 0.9919 | Val Loss: 1.1766 | Val Acc: 0.7444\n",
      "Epoch 15/50 | Train Loss: 0.0338 | Train Acc: 0.9898 | Val Loss: 1.2173 | Val Acc: 0.7263\n",
      "Epoch 16/50 | Train Loss: 0.0246 | Train Acc: 0.9919 | Val Loss: 1.2698 | Val Acc: 0.7249\n",
      "Epoch 17/50 | Train Loss: 0.0151 | Train Acc: 0.9961 | Val Loss: 1.2609 | Val Acc: 0.7374\n",
      "Epoch 18/50 | Train Loss: 0.0322 | Train Acc: 0.9925 | Val Loss: 1.4030 | Val Acc: 0.7123\n",
      "Epoch 19/50 | Train Loss: 0.0502 | Train Acc: 0.9862 | Val Loss: 1.4165 | Val Acc: 0.6983\n",
      "Epoch 20/50 | Train Loss: 0.0536 | Train Acc: 0.9833 | Val Loss: 1.3590 | Val Acc: 0.7249\n",
      "Epoch 21/50 | Train Loss: 0.0279 | Train Acc: 0.9922 | Val Loss: 1.3154 | Val Acc: 0.7235\n",
      "Epoch 22/50 | Train Loss: 0.0454 | Train Acc: 0.9856 | Val Loss: 1.9186 | Val Acc: 0.6494\n",
      "Epoch 23/50 | Train Loss: 0.0295 | Train Acc: 0.9916 | Val Loss: 1.2821 | Val Acc: 0.7304\n",
      "Epoch 24/50 | Train Loss: 0.0140 | Train Acc: 0.9955 | Val Loss: 1.8883 | Val Acc: 0.6858\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[19], line 22\u001b[0m\n\u001b[1;32m     19\u001b[0m loss\u001b[38;5;241m.\u001b[39mbackward()\n\u001b[1;32m     20\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mstep()\n\u001b[0;32m---> 22\u001b[0m train_loss \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[43mloss\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitem\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;241m*\u001b[39m images\u001b[38;5;241m.\u001b[39msize(\u001b[38;5;241m0\u001b[39m)\n\u001b[1;32m     23\u001b[0m _, preds \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mmax(outputs, \u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m     24\u001b[0m train_corrects \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39msum(preds \u001b[38;5;241m==\u001b[39m labels)\u001b[38;5;241m.\u001b[39mitem()\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "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",
    "        labels = torch.squeeze(labels)\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",
    "    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,
   "id": "abfc7e69-b918-4b0f-b603-09e651b8c1ea",
   "metadata": {},
   "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,
   "id": "d5960637-8755-4074-8c55-74f2b52f7e1c",
   "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",
   "id": "0334eba0-bfe7-4e8d-a03a-737f96f12cdd",
   "metadata": {},
   "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.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}