ConvNext.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/alimomennasab/ChestXRay-Classification/blob/main/ConvNext.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pP9Y0fbLzw9I"
      },
      "source": [
        "##1. Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nQRvCSqOiaqg"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torchvision.datasets as datasets\n",
        "import torchvision.transforms as transforms\n",
        "import pandas as pd\n",
        "import os\n",
        "import numpy as np\n",
        "from torchvision.ops import StochasticDepth\n",
        "from typing import List\n",
        "from PIL import Image\n",
        "from glob import glob\n",
        "from torch import Tensor\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torchvision.datasets import ImageFolder\n",
        "from torch.optim.lr_scheduler import ReduceLROnPlateau"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7j7RpSAQz34l",
        "outputId": "35655de8-974f-4765-a263-55d725b15a97"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uKRnJ5C-z4Pu"
      },
      "outputs": [],
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CPls5vz6zznp"
      },
      "source": [
        "##2. Loading Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lW6oUmH5SGPs"
      },
      "outputs": [],
      "source": [
        "data_dir = '/content/drive/My Drive/chest_xray/'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HKpTb3pLK1mF"
      },
      "outputs": [],
      "source": [
        "# Split dataset into training, validation, and test sets\n",
        "train_dir = os.path.join(data_dir, 'train')\n",
        "val_dir = os.path.join(data_dir, 'val')\n",
        "test_dir = os.path.join(data_dir, 'test')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zxi0-PleSQQC",
        "outputId": "9b1a5a5d-fa74-420c-889b-72e081c2405a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['test', 'val', 'train']\n",
            "['PNEUMONIA', 'NORMAL']\n"
          ]
        }
      ],
      "source": [
        "print(os.listdir(data_dir))\n",
        "classes = os.listdir(data_dir + \"/train\")\n",
        "print(classes)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Spwu3fZfSWfH",
        "outputId": "07866d07-d208-4d3c-b509-885a8f55140d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "No. of training examples for Pneumonia: 3875\n",
            "['person557_virus_1097.jpeg', 'person553_bacteria_2316.jpeg', 'person537_bacteria_2264.jpeg', 'person543_bacteria_2281.jpeg', 'person496_virus_1003.jpeg']\n"
          ]
        }
      ],
      "source": [
        "# Pneumonia images\n",
        "pneumonia_files = os.listdir(data_dir + \"/train/PNEUMONIA\")\n",
        "print('No. of training examples for Pneumonia:', len(pneumonia_files))\n",
        "print(pneumonia_files[:5])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "k41O-OQZSdOy",
        "outputId": "12997a44-7693-49e6-e138-0a0c23142e9e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "No. of training examples for Normal: 1341\n",
            "['IM-0524-0001.jpeg', 'IM-0515-0001.jpeg', 'IM-0508-0001.jpeg', 'IM-0516-0001.jpeg', 'IM-0511-0001-0002.jpeg']\n"
          ]
        }
      ],
      "source": [
        "# Normal (healthy) images\n",
        "normal_files = os.listdir(data_dir + \"/train/NORMAL\")\n",
        "print('No. of training examples for Normal:', len(normal_files))\n",
        "print(normal_files[:5])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W7f1dkPIyIAM"
      },
      "outputs": [],
      "source": [
        "# There are almost three times more pneumonia images than normal images, so we will use class weighing\n",
        "\n",
        "# Define classes\n",
        "classes = ['NORMAL', 'PNEUMONIA']\n",
        "\n",
        "# Define class weights\n",
        "num_pneumonia_train = len(os.listdir(os.path.join(train_dir, classes[1])))\n",
        "num_normal_train = len(os.listdir(os.path.join(train_dir, classes[0])))\n",
        "total_train = num_pneumonia_train + num_normal_train\n",
        "class_weights = torch.tensor([total_train/num_normal_train, total_train/num_pneumonia_train]).to(device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mW2QV9xoC2oO"
      },
      "source": [
        "##3. Preparing Dataset and DataLoader"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gCrAW01bDOzn"
      },
      "outputs": [],
      "source": [
        "# Define transforms\n",
        "train_transform = transforms.Compose([\n",
        "    transforms.RandomResizedCrop(224),\n",
        "    transforms.RandomHorizontalFlip(),\n",
        "    transforms.RandomAffine(degrees=10, translate=(0.05,0.05)),\n",
        "    transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1),\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize((0.5), (0.5))\n",
        "])\n",
        "\n",
        "val_and_test_transform = transforms.Compose([\n",
        "    transforms.Resize(256),\n",
        "    transforms.CenterCrop(224),\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize((0.5), (0.5))\n",
        "])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GiSPsuL1VNGE"
      },
      "outputs": [],
      "source": [
        "# Define data directories. We won't use a custom class because the dataset is already well-formatted.\n",
        "\n",
        "train_dataset = ImageFolder('/content/drive/My Drive/chest_xray/train', transform = train_transform)\n",
        "val_dataset = ImageFolder('/content/drive/My Drive/chest_xray/val', transform = val_and_test_transform)\n",
        "test_dataset = ImageFolder('/content/drive/My Drive/chest_xray/test', transform = val_and_test_transform)\n",
        "\n",
        "train_loader = DataLoader(train_dataset, batch_size = 16, shuffle = True)\n",
        "val_loader = DataLoader(val_dataset, batch_size = 16, shuffle = False)\n",
        "test_loader = DataLoader(test_dataset, batch_size = 16, shuffle = False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QQuODp6wPduz"
      },
      "source": [
        "##4. Defining and Choosing Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IRL4xlErPfaH"
      },
      "outputs": [],
      "source": [
        "class LayerScaler(nn.Module):\n",
        "    def __init__(self, init_value: float, dimensions: int):\n",
        "        super().__init__()\n",
        "        self.gamma = nn.Parameter(init_value * torch.ones((dimensions)), \n",
        "                                    requires_grad=True)\n",
        "        \n",
        "    def forward(self, x):\n",
        "        return self.gamma[None,...,None,None] * x\n",
        "        \n",
        "class ConvNormAct(nn.Sequential):\n",
        "    # A little util layer composed by (conv) -> (norm) -> (act) layers. \n",
        "    def __init__(\n",
        "        self,\n",
        "        in_features: int,\n",
        "        out_features: int,\n",
        "        kernel_size: int,\n",
        "        norm = nn.BatchNorm2d,\n",
        "        act = nn.ReLU,\n",
        "        **kwargs\n",
        "    ):\n",
        "        super().__init__(\n",
        "            nn.Conv2d(\n",
        "                in_features,\n",
        "                out_features,\n",
        "                kernel_size=kernel_size,\n",
        "                padding=kernel_size // 2,\n",
        "                **kwargs\n",
        "            ),\n",
        "            norm(out_features),\n",
        "            act(),\n",
        "        )\n",
        "\n",
        "class BottleNeckBlock(nn.Module):\n",
        "    def __init__(\n",
        "        self,\n",
        "        in_features: int,\n",
        "        out_features: int,\n",
        "        expansion: int = 4,\n",
        "        drop_p: float = .0,\n",
        "        layer_scaler_init_value: float = 1e-6,\n",
        "    ):\n",
        "        super().__init__()\n",
        "        expanded_features = out_features * expansion\n",
        "        self.block = nn.Sequential(\n",
        "            # narrow -> wide (with depth-wise and bigger kernel)\n",
        "            nn.Conv2d(\n",
        "                in_features, in_features, kernel_size=7, padding=3, bias=False, groups=in_features\n",
        "            ),\n",
        "            # GroupNorm with num_groups=1 is the same as LayerNorm but works for 2D data\n",
        "            nn.GroupNorm(num_groups=1, num_channels=in_features),\n",
        "            # wide -> wide \n",
        "            nn.Conv2d(in_features, expanded_features, kernel_size=1),\n",
        "            nn.GELU(),\n",
        "            # wide -> narrow\n",
        "            nn.Conv2d(expanded_features, out_features, kernel_size=1),\n",
        "        )\n",
        "        self.layer_scaler = LayerScaler(layer_scaler_init_value, out_features)\n",
        "        self.drop_path = StochasticDepth(drop_p, mode=\"batch\")\n",
        "\n",
        "        \n",
        "    def forward(self, x: Tensor) -> Tensor:\n",
        "        res = x\n",
        "        x = self.block(x)\n",
        "        x = self.layer_scaler(x)\n",
        "        x = self.drop_path(x)\n",
        "        x += res\n",
        "        return x\n",
        "\n",
        "class ConvNexStage(nn.Sequential):\n",
        "    def __init__(\n",
        "        self, in_features: int, out_features: int, depth: int, **kwargs\n",
        "    ):\n",
        "        super().__init__(\n",
        "            # add the downsampler\n",
        "            nn.Sequential(\n",
        "                nn.GroupNorm(num_groups=1, num_channels=in_features),\n",
        "                nn.Conv2d(in_features, out_features, kernel_size=2, stride=2)\n",
        "            ),\n",
        "            *[\n",
        "                BottleNeckBlock(out_features, out_features, **kwargs)\n",
        "                for _ in range(depth)\n",
        "            ],\n",
        "        )\n",
        "\n",
        "class ConvNextStem(nn.Sequential):\n",
        "    def __init__(self, in_features: int, out_features: int):\n",
        "        super().__init__(\n",
        "            nn.Conv2d(in_features, out_features, kernel_size=4, stride=4),\n",
        "            nn.BatchNorm2d(out_features)\n",
        "        )\n",
        "\n",
        "class ConvNextEncoder(nn.Module):\n",
        "    def __init__(\n",
        "        self,\n",
        "        in_channels: int,\n",
        "        stem_features: int,\n",
        "        depths: List[int],\n",
        "        widths: List[int],\n",
        "        drop_p: float = .0,\n",
        "    ):\n",
        "        super().__init__()\n",
        "        self.stem = ConvNextStem(in_channels, stem_features)\n",
        "\n",
        "        in_out_widths = list(zip(widths, widths[1:]))\n",
        "        # create drop paths probabilities (one for each stage)\n",
        "        drop_probs = [x.item() for x in torch.linspace(0, drop_p, sum(depths))] \n",
        "        \n",
        "        self.stages = nn.ModuleList(\n",
        "            [\n",
        "                ConvNexStage(stem_features, widths[0], depths[0], drop_p=drop_probs[0]),\n",
        "                *[\n",
        "                    ConvNexStage(in_features, out_features, depth, drop_p=drop_p)\n",
        "                    for (in_features, out_features), depth, drop_p in zip(\n",
        "                        in_out_widths, depths[1:], drop_probs[1:]\n",
        "                    )\n",
        "                ],\n",
        "            ]\n",
        "        )\n",
        "        \n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.stem(x)\n",
        "        for stage in self.stages:\n",
        "            x = stage(x)\n",
        "        return x\n",
        "\n",
        "class ClassificationHead(nn.Sequential):\n",
        "    def __init__(self, num_channels: int, num_classes: int = 2):\n",
        "        super().__init__(\n",
        "            nn.AdaptiveAvgPool2d((1, 1)),\n",
        "            nn.Flatten(1),\n",
        "            nn.LayerNorm(num_channels),\n",
        "            nn.Linear(num_channels, num_classes)\n",
        "        )\n",
        "    \n",
        "class ConvNext(nn.Sequential):\n",
        "    def __init__(self,  \n",
        "                 in_channels: int,\n",
        "                 stem_features: int,\n",
        "                 depths: List[int],\n",
        "                 widths: List[int],\n",
        "                 drop_p: float = .0,\n",
        "                 num_classes: int = 2):\n",
        "        super().__init__()\n",
        "        self.encoder = ConvNextEncoder(in_channels, stem_features, depths, widths, drop_p)\n",
        "        self.head = ClassificationHead(widths[-1], num_classes)\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#Initializing model\n",
        "model = ConvNext(in_channels=3, stem_features=64, depths=[3,4,6,4], widths=[256, 512, 1024, 2048])"
      ],
      "metadata": {
        "id": "r9VmuzyS9r6Y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "if torch.cuda.is_available():\n",
        "    model.cuda()"
      ],
      "metadata": {
        "id": "NQ3GBuIl9sdg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f1dyk_TGPu73"
      },
      "source": [
        "##5. Testing Model"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x = torch.rand(1, 3, 224, 224).to(device)\n",
        "model(x).shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TWmF4W_BnskO",
        "outputId": "f7a0d8c4-d15d-4123-e358-65bcb3de5c13"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([1, 2])"
            ]
          },
          "metadata": {},
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4Ts-Y_GmP1d2"
      },
      "source": [
        "##6. Defining Main Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OFI29VFCP5EE"
      },
      "outputs": [],
      "source": [
        "# Hyper-parameters\n",
        "num_epochs = 30\n",
        "learning_rate = 0.001\n",
        "patience = 10"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "v0PQts-aP5f5",
        "outputId": "d60845ec-bd95-45a1-dd71-14947666fc27"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch [1/30], Step [100/326] Loss: 0.8381\n",
            "Epoch [1/30], Step [200/326] Loss: 0.4606\n",
            "Epoch [1/30], Step [300/326] Loss: 0.8287\n",
            "Saving model with validation loss of 0.5096...\n",
            "Epoch [2/30], Step [100/326] Loss: 0.5864\n",
            "Epoch [2/30], Step [200/326] Loss: 0.4673\n",
            "Epoch [2/30], Step [300/326] Loss: 0.5147\n",
            "Saving model with validation loss of 0.4660...\n",
            "Epoch [3/30], Step [100/326] Loss: 0.5141\n",
            "Epoch [3/30], Step [200/326] Loss: 0.5784\n",
            "Epoch [3/30], Step [300/326] Loss: 0.4019\n",
            "Epoch [4/30], Step [100/326] Loss: 0.6047\n",
            "Epoch [4/30], Step [200/326] Loss: 0.3296\n",
            "Epoch [4/30], Step [300/326] Loss: 0.4036\n",
            "Saving model with validation loss of 0.4536...\n",
            "Epoch [5/30], Step [100/326] Loss: 0.7433\n",
            "Epoch [5/30], Step [200/326] Loss: 0.4903\n",
            "Epoch [5/30], Step [300/326] Loss: 0.3827\n",
            "Epoch [6/30], Step [100/326] Loss: 0.4108\n",
            "Epoch [6/30], Step [200/326] Loss: 0.7480\n",
            "Epoch [6/30], Step [300/326] Loss: 0.5949\n",
            "Epoch [7/30], Step [100/326] Loss: 0.5605\n",
            "Epoch [7/30], Step [200/326] Loss: 0.5370\n",
            "Epoch [7/30], Step [300/326] Loss: 0.6100\n",
            "Epoch [8/30], Step [100/326] Loss: 0.4821\n",
            "Epoch [8/30], Step [200/326] Loss: 0.5782\n",
            "Epoch [8/30], Step [300/326] Loss: 0.4933\n",
            "Epoch [9/30], Step [100/326] Loss: 0.4274\n",
            "Epoch [9/30], Step [200/326] Loss: 0.6022\n",
            "Epoch [9/30], Step [300/326] Loss: 0.3393\n",
            "Epoch [10/30], Step [100/326] Loss: 0.2752\n",
            "Epoch [10/30], Step [200/326] Loss: 0.3951\n",
            "Epoch [10/30], Step [300/326] Loss: 0.3502\n",
            "Epoch [11/30], Step [100/326] Loss: 0.7066\n",
            "Epoch [11/30], Step [200/326] Loss: 0.5653\n",
            "Epoch [11/30], Step [300/326] Loss: 0.5031\n",
            "Epoch [12/30], Step [100/326] Loss: 0.2378\n",
            "Epoch [12/30], Step [200/326] Loss: 0.3944\n",
            "Epoch [12/30], Step [300/326] Loss: 0.3838\n",
            "Epoch [13/30], Step [100/326] Loss: 0.5490\n",
            "Epoch [13/30], Step [200/326] Loss: 0.4318\n",
            "Epoch [13/30], Step [300/326] Loss: 0.2483\n",
            "Epoch [14/30], Step [100/326] Loss: 0.4577\n",
            "Epoch [14/30], Step [200/326] Loss: 0.4144\n",
            "Epoch [14/30], Step [300/326] Loss: 0.4653\n",
            "Epoch [15/30], Step [100/326] Loss: 0.2947\n",
            "Epoch [15/30], Step [200/326] Loss: 0.5949\n",
            "Epoch [15/30], Step [300/326] Loss: 0.2707\n",
            "Epoch [16/30], Step [100/326] Loss: 0.8715\n",
            "Epoch [16/30], Step [200/326] Loss: 0.6582\n",
            "Epoch [16/30], Step [300/326] Loss: 0.4430\n",
            "Epoch [17/30], Step [100/326] Loss: 0.5877\n",
            "Epoch [17/30], Step [200/326] Loss: 0.6500\n",
            "Epoch [17/30], Step [300/326] Loss: 0.3656\n",
            "Saving model with validation loss of 0.4177...\n",
            "Epoch [18/30], Step [100/326] Loss: 0.5408\n",
            "Epoch [18/30], Step [200/326] Loss: 0.3458\n",
            "Epoch [18/30], Step [300/326] Loss: 0.2899\n",
            "Epoch [19/30], Step [100/326] Loss: 0.2060\n",
            "Epoch [19/30], Step [200/326] Loss: 0.4114\n",
            "Epoch [19/30], Step [300/326] Loss: 0.4726\n",
            "Epoch [20/30], Step [100/326] Loss: 0.2660\n",
            "Epoch [20/30], Step [200/326] Loss: 0.5408\n",
            "Epoch [20/30], Step [300/326] Loss: 0.6458\n",
            "Epoch [21/30], Step [100/326] Loss: 0.2767\n",
            "Epoch [21/30], Step [200/326] Loss: 0.2291\n",
            "Epoch [21/30], Step [300/326] Loss: 0.3474\n",
            "Epoch [22/30], Step [100/326] Loss: 0.3798\n",
            "Epoch [22/30], Step [200/326] Loss: 0.6556\n",
            "Epoch [22/30], Step [300/326] Loss: 0.1407\n",
            "Epoch [23/30], Step [100/326] Loss: 0.3449\n",
            "Epoch [23/30], Step [200/326] Loss: 0.4029\n",
            "Epoch [23/30], Step [300/326] Loss: 0.5506\n",
            "Epoch [24/30], Step [100/326] Loss: 0.2209\n",
            "Epoch [24/30], Step [200/326] Loss: 0.3019\n",
            "Epoch [24/30], Step [300/326] Loss: 1.0342\n",
            "Epoch [25/30], Step [100/326] Loss: 0.3958\n",
            "Epoch [25/30], Step [200/326] Loss: 0.5548\n",
            "Epoch [25/30], Step [300/326] Loss: 0.3271\n",
            "Epoch [26/30], Step [100/326] Loss: 0.4096\n",
            "Epoch [26/30], Step [200/326] Loss: 0.4127\n",
            "Epoch [26/30], Step [300/326] Loss: 0.6600\n",
            "Epoch [27/30], Step [100/326] Loss: 0.5709\n",
            "Epoch [27/30], Step [200/326] Loss: 0.2892\n",
            "Epoch [27/30], Step [300/326] Loss: 0.3226\n",
            "Epoch [28/30], Step [100/326] Loss: 0.6224\n",
            "Epoch [28/30], Step [200/326] Loss: 0.2720\n",
            "Epoch [28/30], Step [300/326] Loss: 0.5048\n",
            "Epoch [29/30], Step [100/326] Loss: 0.4290\n",
            "Epoch [29/30], Step [200/326] Loss: 0.2694\n",
            "Epoch [29/30], Step [300/326] Loss: 0.4416\n",
            "Epoch [30/30], Step [100/326] Loss: 0.4790\n",
            "Epoch [30/30], Step [200/326] Loss: 0.1100\n",
            "Epoch [30/30], Step [300/326] Loss: 0.5521\n"
          ]
        }
      ],
      "source": [
        "# Loss and optimizer\n",
        "criterion = nn.CrossEntropyLoss(weight=class_weights)\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
        "\n",
        "curr_lr = learning_rate\n",
        "total_step = len(train_loader)\n",
        "\n",
        "# Early stopping parameters\n",
        "early_stopping_counter = 0\n",
        "best_loss = float('inf')\n",
        "\n",
        "# For updating learning rate\n",
        "def update_lr(optimizer, lr):    \n",
        "    for param_group in optimizer.param_groups:\n",
        "        param_group['lr'] = lr\n",
        "\n",
        "\n",
        "for epoch in range(num_epochs):\n",
        "    for i, (images, labels) in enumerate(train_loader):\n",
        "        images, labels = images.to(device), labels.to(device)\n",
        "        \n",
        "        # Forward pass\n",
        "        outputs = model(images)\n",
        "        loss = criterion(outputs, labels)\n",
        "        \n",
        "        # Backward and optimize\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        \n",
        "        if (i+1) % 100 == 0:\n",
        "            print (\"Epoch [{}/{}], Step [{}/{}] Loss: {:.4f}\"\n",
        "                   .format(epoch+1, num_epochs, i+1, total_step, loss.item()))\n",
        "    \n",
        "    # Decay learning rate\n",
        "    if (epoch+1) % 20 == 0:\n",
        "        curr_lr /= 3\n",
        "        update_lr(optimizer, curr_lr)\n",
        "\n",
        "    # Calculate validation loss\n",
        "    with torch.no_grad():\n",
        "        val_loss = 0\n",
        "        for images, labels in val_loader:\n",
        "            images, labels = images.to(device), labels.to(device)\n",
        "            outputs = model(images)\n",
        "            val_loss += criterion(outputs, labels)\n",
        "        val_loss /= len(val_loader)\n",
        "\n",
        "        # Save the model if the validation loss is the best observed yet\n",
        "        if val_loss < best_loss:\n",
        "            print(f'Saving model with validation loss of {val_loss:.4f}...')\n",
        "            torch.save(model.state_dict(), 'best_model.pth')\n",
        "            best_loss = val_loss\n",
        "\n",
        "    # Early stopping if overfitting\n",
        "    if early_stopping_counter >= patience:\n",
        "        print(f'Validation loss has not improved for {patience} epochs. Early stopping...')\n",
        "        break\n",
        "    elif val_loss < best_loss:\n",
        "        best_loss = val_loss\n",
        "        early_stopping_counter = 0\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PXmE1g9WQ6_2"
      },
      "source": [
        "##7. Testing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yW_Iy6SqQ8qE",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "98a15e49-1170-4caa-a615-b35b01c0a773"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model accuracy on test images: 70.99358974358974 %\n"
          ]
        }
      ],
      "source": [
        "# Load the saved model checkpoint\n",
        "checkpoint = torch.load('best_model.pth')\n",
        "model.load_state_dict(checkpoint)\n",
        "\n",
        "model.eval()\n",
        "with torch.no_grad():\n",
        "    correct = 0\n",
        "    total = 0\n",
        "    for images, labels in test_loader:\n",
        "        images = images.to(device)\n",
        "        labels = labels.to(device)\n",
        "        outputs = model(images)\n",
        "        _, predicted = torch.max(outputs.data, 1)\n",
        "        total += labels.size(0)\n",
        "        correct += (predicted == labels).sum().item()\n",
        "\n",
        "    print('Model accuracy on test images: {} %'.format(100 * correct / total))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}