ResNext.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/ResNext.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Notes: ResNext is combo of VGGNet, ResNet, and InceptionNet"
      ],
      "metadata": {
        "id": "A4A3jeva7gaJ"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pP9Y0fbLzw9I"
      },
      "source": [
        "##1. Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "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 torch.nn.functional as F\n",
        "import pandas as pd\n",
        "import math\n",
        "import os\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "from glob import glob\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": 45,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7j7RpSAQz34l",
        "outputId": "cf80f7cc-45f4-4bbe-b966-209633994b24"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 46,
      "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": 47,
      "metadata": {
        "id": "lW6oUmH5SGPs"
      },
      "outputs": [],
      "source": [
        "data_dir = '/content/drive/My Drive/chest_xray/'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 48,
      "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": 49,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zxi0-PleSQQC",
        "outputId": "48413fe4-416b-4c29-c929-bb8b95718dfe"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['train', 'chest_xray', 'val', '__MACOSX', 'test']\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": 50,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Spwu3fZfSWfH",
        "outputId": "b4f026b3-3714-4368-cdb5-7cc1db834f2c"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "No. of training examples for Pneumonia: 3875\n",
            "['person556_virus_1096.jpeg', 'person536_bacteria_2257.jpeg', 'person581_bacteria_2390.jpeg', 'person592_bacteria_2434.jpeg', 'person581_virus_1125.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": 51,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "k41O-OQZSdOy",
        "outputId": "fc82b2ea-6153-470d-a844-08e00ba8544b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "No. of training examples for Normal: 1341\n",
            "['IM-0526-0001.jpeg', 'IM-0524-0001.jpeg', 'IM-0507-0001.jpeg', 'IM-0508-0001.jpeg', 'IM-0520-0001.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": 52,
      "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": 53,
      "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": 54,
      "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": 55,
      "metadata": {
        "id": "IRL4xlErPfaH"
      },
      "outputs": [],
      "source": [
        "class Bottleneck(nn.Module):\n",
        "\n",
        "    expansion = 4\n",
        "\n",
        "    def __init__(self, inplanes, planes, baseWidth, cardinality, stride=1, downsample=None):\n",
        "      \n",
        "        # cardinality: num of convolution groups.\n",
        "\n",
        "        super(Bottleneck, self).__init__()\n",
        "\n",
        "        D = int(math.floor(planes * (baseWidth / 64)))\n",
        "        C = cardinality\n",
        "\n",
        "        self.conv1 = nn.Conv2d(inplanes, D*C, kernel_size=1, stride=1, padding=0, bias=False)\n",
        "        self.bn1 = nn.BatchNorm2d(D*C)\n",
        "        self.conv2 = nn.Conv2d(D*C, D*C, kernel_size=3, stride=stride, padding=1, groups=C, bias=False)\n",
        "        self.bn2 = nn.BatchNorm2d(D*C)\n",
        "        self.conv3 = nn.Conv2d(D*C, planes * 4, kernel_size=1, stride=1, padding=0, bias=False)\n",
        "        self.bn3 = nn.BatchNorm2d(planes * 4)\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "\n",
        "        self.downsample = downsample\n",
        "\n",
        "    def forward(self, x):\n",
        "        residual = x\n",
        "\n",
        "        out = self.conv1(x)\n",
        "        out = self.bn1(out)\n",
        "        out = self.relu(out)\n",
        "\n",
        "        out = self.conv2(out)\n",
        "        out = self.bn2(out)\n",
        "        out = self.relu(out)\n",
        "\n",
        "        out = self.conv3(out)\n",
        "        out = self.bn3(out)\n",
        "\n",
        "        if self.downsample is not None:\n",
        "            residual = self.downsample(x)\n",
        "\n",
        "        out += residual\n",
        "        out = self.relu(out)\n",
        "\n",
        "        return out\n",
        "\n",
        "\n",
        "class ResNeXt(nn.Module):\n",
        "    def __init__(self, baseWidth, cardinality, layers, num_classes):\n",
        "\n",
        "        super(ResNeXt, self).__init__()\n",
        "        block = Bottleneck\n",
        "\n",
        "        self.cardinality = cardinality\n",
        "        self.baseWidth = baseWidth\n",
        "        self.num_classes = num_classes\n",
        "        self.inplanes = 64\n",
        "        self.output_size = 64\n",
        "\n",
        "        self.conv1 = nn.Conv2d(3, 64, 7, 2, 3, bias=False)\n",
        "        self.bn1 = nn.BatchNorm2d(64)\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "        self.maxpool1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
        "        self.layer1 = self._make_layer(block, 64, layers[0])\n",
        "        self.layer2 = self._make_layer(block, 128, layers[1], 2)\n",
        "        self.layer3 = self._make_layer(block, 256, layers[2], 2)\n",
        "        self.layer4 = self._make_layer(block, 512, layers[3], 2)\n",
        "        self.avgpool = nn.AvgPool2d(7)      \n",
        "        self.fc = nn.Linear(512 * block.expansion, num_classes)\n",
        "\n",
        "        for m in self.modules():\n",
        "            if isinstance(m, nn.Conv2d):\n",
        "                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n",
        "                m.weight.data.normal_(0, math.sqrt(2. / n))\n",
        "            elif isinstance(m, nn.BatchNorm2d):\n",
        "                m.weight.data.fill_(1)\n",
        "                m.bias.data.zero_()\n",
        "\n",
        "    def _make_layer(self, block, planes, blocks, stride=1):\n",
        "        downsample = None\n",
        "        if stride != 1 or self.inplanes != planes * block.expansion:\n",
        "            downsample = nn.Sequential(\n",
        "                nn.Conv2d(self.inplanes, planes * block.expansion,\n",
        "                          kernel_size=1, stride=stride, bias=False),\n",
        "                nn.BatchNorm2d(planes * block.expansion),\n",
        "            )\n",
        "\n",
        "        layers = []\n",
        "        layers.append(block(self.inplanes, planes, self.baseWidth, self.cardinality, stride, downsample))\n",
        "        self.inplanes = planes * block.expansion\n",
        "        for i in range(1, blocks):\n",
        "            layers.append(block(self.inplanes, planes, self.baseWidth, self.cardinality))\n",
        "\n",
        "        return nn.Sequential(*layers)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.conv1(x)\n",
        "        x = self.bn1(x)\n",
        "        x = self.relu(x)\n",
        "        x = self.maxpool1(x)\n",
        "        x = self.layer1(x)\n",
        "        x = self.layer2(x)\n",
        "        x = self.layer3(x)\n",
        "        x = self.layer4(x)\n",
        "        x = self.avgpool(x)\n",
        "        x = x.view(x.size(0), -1)\n",
        "        x = self.fc(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "# 2 classes\n",
        "def ResNext50(baseWidth, cardinality):\n",
        "    model = ResNeXt(baseWidth, cardinality, [3, 4, 6, 3], 2)\n",
        "    return model\n",
        "\n",
        "\n",
        "def ResNext101(baseWidth, cardinality):\n",
        "    model = ResNeXt(baseWidth, cardinality, [3, 4, 23, 3], 2)\n",
        "    return model\n",
        "\n",
        "\n",
        "def ResNext152(baseWidth, cardinality):\n",
        "    model = ResNeXt(baseWidth, cardinality, [3, 8, 36, 3], 2)\n",
        "    return model"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Initializing model\n",
        "model = ResNext152(baseWidth = 4, cardinality = 32)"
      ],
      "metadata": {
        "id": "r9VmuzyS9r6Y"
      },
      "execution_count": 56,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "if torch.cuda.is_available():\n",
        "    model.cuda()"
      ],
      "metadata": {
        "id": "NQ3GBuIl9sdg"
      },
      "execution_count": 57,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f1dyk_TGPu73"
      },
      "source": [
        "##5. Testing Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 58,
      "metadata": {
        "id": "QcsXcSO3PuA7",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f59e3221-6f4d-4655-dbdb-d39a5ebb5b7b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([8, 2])\n"
          ]
        }
      ],
      "source": [
        "x = torch.randn(8, 3, 224, 224).to(device)\n",
        "print(model(x).shape) "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4Ts-Y_GmP1d2"
      },
      "source": [
        "##6. Defining Main Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 59,
      "metadata": {
        "id": "OFI29VFCP5EE"
      },
      "outputs": [],
      "source": [
        "# Hyper-parameters\n",
        "num_epochs = 30\n",
        "learning_rate = 0.001\n",
        "patience = 10"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 60,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "v0PQts-aP5f5",
        "outputId": "d0dbed5e-82d5-4acd-9fac-4fb2c0229b28"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch [1/30], Step [100/326] Loss: 0.7175\n",
            "Epoch [1/30], Step [200/326] Loss: 0.7366\n",
            "Epoch [1/30], Step [300/326] Loss: 0.5713\n",
            "Saving model with validation loss of 0.6458...\n",
            "Epoch [2/30], Step [100/326] Loss: 0.5254\n",
            "Epoch [2/30], Step [200/326] Loss: 0.4545\n",
            "Epoch [2/30], Step [300/326] Loss: 0.4579\n",
            "Epoch [3/30], Step [100/326] Loss: 0.6595\n",
            "Epoch [3/30], Step [200/326] Loss: 0.5817\n",
            "Epoch [3/30], Step [300/326] Loss: 0.5827\n",
            "Epoch [4/30], Step [100/326] Loss: 0.5870\n",
            "Epoch [4/30], Step [200/326] Loss: 0.6499\n",
            "Epoch [4/30], Step [300/326] Loss: 0.5449\n",
            "Epoch [5/30], Step [100/326] Loss: 0.4706\n",
            "Epoch [5/30], Step [200/326] Loss: 0.3053\n",
            "Epoch [5/30], Step [300/326] Loss: 0.4233\n",
            "Epoch [6/30], Step [100/326] Loss: 0.4812\n",
            "Epoch [6/30], Step [200/326] Loss: 0.5639\n",
            "Epoch [6/30], Step [300/326] Loss: 0.2711\n",
            "Epoch [7/30], Step [100/326] Loss: 0.6003\n",
            "Epoch [7/30], Step [200/326] Loss: 0.6206\n",
            "Epoch [7/30], Step [300/326] Loss: 0.3373\n",
            "Epoch [8/30], Step [100/326] Loss: 0.5192\n",
            "Epoch [8/30], Step [200/326] Loss: 0.5139\n",
            "Epoch [8/30], Step [300/326] Loss: 0.4855\n",
            "Epoch [9/30], Step [100/326] Loss: 0.5520\n",
            "Epoch [9/30], Step [200/326] Loss: 0.5076\n",
            "Epoch [9/30], Step [300/326] Loss: 0.4743\n",
            "Epoch [10/30], Step [100/326] Loss: 0.3995\n",
            "Epoch [10/30], Step [200/326] Loss: 0.4003\n",
            "Epoch [10/30], Step [300/326] Loss: 0.3245\n",
            "Epoch [11/30], Step [100/326] Loss: 0.8181\n",
            "Epoch [11/30], Step [200/326] Loss: 0.4125\n",
            "Epoch [11/30], Step [300/326] Loss: 0.4365\n",
            "Epoch [12/30], Step [100/326] Loss: 0.2290\n",
            "Epoch [12/30], Step [200/326] Loss: 0.5196\n",
            "Epoch [12/30], Step [300/326] Loss: 0.2967\n",
            "Epoch [13/30], Step [100/326] Loss: 0.4197\n",
            "Epoch [13/30], Step [200/326] Loss: 0.2796\n",
            "Epoch [13/30], Step [300/326] Loss: 0.6188\n",
            "Epoch [14/30], Step [100/326] Loss: 0.1507\n",
            "Epoch [14/30], Step [200/326] Loss: 0.3876\n",
            "Epoch [14/30], Step [300/326] Loss: 0.2921\n",
            "Epoch [15/30], Step [100/326] Loss: 0.6215\n",
            "Epoch [15/30], Step [200/326] Loss: 0.2397\n",
            "Epoch [15/30], Step [300/326] Loss: 0.3003\n",
            "Epoch [16/30], Step [100/326] Loss: 0.5310\n",
            "Epoch [16/30], Step [200/326] Loss: 0.2775\n",
            "Epoch [16/30], Step [300/326] Loss: 0.1563\n",
            "Epoch [17/30], Step [100/326] Loss: 0.4474\n",
            "Epoch [17/30], Step [200/326] Loss: 0.3540\n",
            "Epoch [17/30], Step [300/326] Loss: 0.7146\n",
            "Epoch [18/30], Step [100/326] Loss: 0.1610\n",
            "Epoch [18/30], Step [200/326] Loss: 0.9356\n",
            "Epoch [18/30], Step [300/326] Loss: 0.4835\n",
            "Epoch [19/30], Step [100/326] Loss: 0.3936\n",
            "Epoch [19/30], Step [200/326] Loss: 0.2849\n",
            "Epoch [19/30], Step [300/326] Loss: 0.4160\n",
            "Epoch [20/30], Step [100/326] Loss: 0.5205\n",
            "Epoch [20/30], Step [200/326] Loss: 0.4793\n",
            "Epoch [20/30], Step [300/326] Loss: 0.3544\n",
            "Epoch [21/30], Step [100/326] Loss: 0.2255\n",
            "Epoch [21/30], Step [200/326] Loss: 0.2353\n",
            "Epoch [21/30], Step [300/326] Loss: 0.3874\n",
            "Epoch [22/30], Step [100/326] Loss: 0.4347\n",
            "Epoch [22/30], Step [200/326] Loss: 0.1580\n",
            "Epoch [22/30], Step [300/326] Loss: 0.1741\n",
            "Epoch [23/30], Step [100/326] Loss: 0.5107\n",
            "Epoch [23/30], Step [200/326] Loss: 0.1884\n",
            "Epoch [23/30], Step [300/326] Loss: 0.6781\n",
            "Epoch [24/30], Step [100/326] Loss: 0.5510\n",
            "Epoch [24/30], Step [200/326] Loss: 0.3124\n",
            "Epoch [24/30], Step [300/326] Loss: 0.4431\n",
            "Epoch [25/30], Step [100/326] Loss: 0.2886\n",
            "Epoch [25/30], Step [200/326] Loss: 0.2003\n",
            "Epoch [25/30], Step [300/326] Loss: 0.2272\n",
            "Epoch [26/30], Step [100/326] Loss: 0.3070\n",
            "Epoch [26/30], Step [200/326] Loss: 0.4215\n",
            "Epoch [26/30], Step [300/326] Loss: 0.5900\n",
            "Epoch [27/30], Step [100/326] Loss: 0.3516\n",
            "Epoch [27/30], Step [200/326] Loss: 0.1890\n",
            "Epoch [27/30], Step [300/326] Loss: 0.2084\n",
            "Epoch [28/30], Step [100/326] Loss: 0.3420\n",
            "Epoch [28/30], Step [200/326] Loss: 0.4810\n",
            "Epoch [28/30], Step [300/326] Loss: 0.3945\n",
            "Epoch [29/30], Step [100/326] Loss: 0.1462\n",
            "Epoch [29/30], Step [200/326] Loss: 0.2553\n",
            "Epoch [29/30], Step [300/326] Loss: 0.2037\n",
            "Epoch [30/30], Step [100/326] Loss: 0.2209\n",
            "Epoch [30/30], Step [200/326] Loss: 0.1108\n",
            "Epoch [30/30], Step [300/326] Loss: 0.2036\n",
            "Saving model with validation loss of 0.5833...\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": 61,
      "metadata": {
        "id": "yW_Iy6SqQ8qE",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "780afde1-c3a9-4a47-9f51-395ef792ba9c"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model accuracy on test images: 92.46794871794872 %\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": [],
      "authorship_tag": "ABX9TyOgdJDbnVRi19NkAdBNMXUt",
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}