finetune_resnet18.ipynb

{
 "metadata": {
  "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.8.5-final"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import libraries\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "import torchvision\n",
    "import torchvision.datasets as datasets\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data.sampler import SubsetRandomSampler\n",
    "\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from utils import get_num_correct\n",
    "from custom_dataset import ChestXRayDataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "metadata": {},
     "execution_count": 2
    }
   ],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # set the device type\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# declare tranformations\n",
    "transform = {\n",
    "    'train': transforms.Compose([\n",
    "        transforms.Resize((224, 224)),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(\n",
    "            mean=[0.485, 0.456, 0.406],\n",
    "            std=[0.229, 0.224, 0.225]\n",
    "        )\n",
    "    ]),\n",
    "    'test': transforms.Compose([\n",
    "        transforms.Resize((224, 224)),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(\n",
    "            mean=[0.485, 0.456, 0.406],\n",
    "            std=[0.229, 0.224, 0.225]\n",
    "        )\n",
    "    ])\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get train and test directories\n",
    "dirs = {\n",
    "    'train': {\n",
    "        'covid': 'COVID-19 Radiography Database/train/covid',\n",
    "        'normal': 'COVID-19 Radiography Database/train/normal',\n",
    "        'viral': 'COVID-19 Radiography Database/train/viral'\n",
    "    },\n",
    "    'test': 'COVID-19 Radiography Database/test'\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "found 179 covid examples.\nfound 1301 normal examples.\nfound 1305 viral examples.\n"
     ]
    }
   ],
   "source": [
    "# prepare the train data-loader\n",
    "train_set = ChestXRayDataset(dirs['train'], transform['train'])\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size=8, shuffle=True, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare the test and validation data-loader\n",
    "test_set = datasets.ImageFolder(dirs['test'], transform['test'])\n",
    "\n",
    "valid_size = 0.5  # fraction of test_set to be used as validation set\n",
    "\n",
    "# obtain test indices that will be used for validation\n",
    "num_test = len(test_set)\n",
    "indices = list(range(num_test))\n",
    "np.random.shuffle(indices)\n",
    "split = int(np.floor(valid_size*num_test))\n",
    "test_idx, valid_idx = indices[split:], indices[:split]\n",
    "\n",
    "# define samplers for obtaining test and validation batches\n",
    "valid_sampler = SubsetRandomSampler(valid_idx)\n",
    "test_sampler = SubsetRandomSampler(test_idx)\n",
    "\n",
    "# prepare the data loaders\n",
    "valid_loader = torch.utils.data.DataLoader(test_set, batch_size=8, sampler=valid_sampler, num_workers=2)\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size=8, sampler=test_sampler, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "metadata": {},
     "execution_count": 7
    }
   ],
   "source": [
    "# resnet18 = torchvision.models.resnet18(pretrained=True)\n",
    "resnet18 = torchvision.models.resnet18(pretrained=False)\n",
    "resnet18.load_state_dict(\n",
    "    torch.load('../models/resnet18.pth',\n",
    "    map_location=device)\n",
    ")  # load the pretrained resnet18 model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "ResNet(\n",
       "  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (relu): ReLU(inplace=True)\n",
       "  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "  (layer1): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (layer2): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (layer3): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (layer4): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
       "  (fc): Linear(in_features=512, out_features=3, bias=True)\n",
       ")"
      ]
     },
     "metadata": {},
     "execution_count": 8
    }
   ],
   "source": [
    "# change the last fc layer so that it could output 3 classes \n",
    "resnet18.fc = torch.nn.Linear(in_features=512, out_features=3)\n",
    "resnet18.to(device)  # move to GPU (if available)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "Epoch [ 1/3]: 100%|██████████| 349/349 [08:58<00:00,  1.54s/it, acc=0.92, loss=1.31]\n",
      "\t\tAvg training loss: 0.235437\tAvg validation loss: 0.083287\n",
      "\t\tvalid_loss decreased (inf --> 0.083287)  saving model...\n",
      "Epoch [ 2/3]: 100%|██████████| 349/349 [08:45<00:00,  1.51s/it, acc=0.964, loss=0.885]\n",
      "\t\tAvg training loss: 0.111461\tAvg validation loss: 0.097972\n",
      "Epoch [ 3/3]: 100%|██████████| 349/349 [08:38<00:00,  1.49s/it, acc=0.983, loss=1.03]\n",
      "\t\tAvg training loss: 0.061282\tAvg validation loss: 0.061795\n",
      "\t\tvalid_loss decreased (0.083287 --> 0.061795)  saving model...\n"
     ]
    }
   ],
   "source": [
    "criterion = nn.CrossEntropyLoss()  # loss function (categorical cross-entropy)\n",
    "optimizer = optim.Adam(resnet18.parameters(), lr=3e-5)\n",
    "\n",
    "comment = '-resnet18-covid'  # # will be used for naming the run\n",
    "tb = SummaryWriter(comment=comment)\n",
    "\n",
    "# initialize tracker for minimum validation loss\n",
    "valid_loss_min = np.Inf  # set initial minimum to infinity\n",
    "num_epochs = 3  # number of epochs used for training\n",
    "\n",
    "len_train = len(train_set)\n",
    "len_val = len(valid_loader.sampler)\n",
    "len_test = len(test_loader.sampler)\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    train_loss, train_correct = 0, 0  # will be used to track the running loss and correct\n",
    "    #######################\n",
    "    # fine-tune the model #\n",
    "    #######################\n",
    "    train_loop = tqdm(train_loader)\n",
    "    resnet18.train()  # set the model to train mode\n",
    "\n",
    "    for batch in train_loop:\n",
    "        images, labels = batch[0].to(device), batch[1].to(device)  # load the batch to the available device\n",
    "        preds = resnet18(images)  # forward pass\n",
    "        loss = criterion(preds, labels)  # calculate loss\n",
    "        optimizer.zero_grad()  # clear the accumulated gradients from the previous pass\n",
    "        loss.backward()  # backward pass\n",
    "        optimizer.step()  # perform a single optimization step\n",
    "\n",
    "        train_loss += loss.item() * labels.size(0)  # update the running loss\n",
    "        train_correct += get_num_correct(preds, labels)  # update running num correct\n",
    "\n",
    "        train_loop.set_description(f'Epoch [{epoch+1:2d}/{num_epochs}]')\n",
    "        train_loop.set_postfix(loss=loss.item(), acc=train_correct/len_train)\n",
    "\n",
    "    # add train loss and train accuracy for the current epoch to tensorboard\n",
    "    tb.add_scalar('Train Loss', train_loss, epoch)\n",
    "    tb.add_scalar('Train Accuracy', train_correct/len_train, epoch)\n",
    "\n",
    "\n",
    "    resnet18.eval()  # set the model to evaluation mode\n",
    "    with torch.no_grad():  # turn off grad tracking, as we don't need gradients for validation\n",
    "\n",
    "        valid_loss, valid_correct = 0, 0  # will be used to track the running validation loss and correct\n",
    "        ######################\n",
    "        # validate the model #\n",
    "        ######################\n",
    "        for batch in valid_loader:\n",
    "            images, labels = batch[0].to(device), batch[1].to(device)  # load the batch to the available device\n",
    "            preds = resnet18(images)  # forward pass\n",
    "            loss = criterion(preds, labels)  # calculate the loss\n",
    "\n",
    "            valid_loss += loss.item() * labels.size(0)  # update the running loss\n",
    "            valid_correct += get_num_correct(preds, labels)  # update running num correct\n",
    "\n",
    "        # add validation loss and validation accuracy for the current epoch to tensorboard\n",
    "        tb.add_scalar('Validation Loss', valid_loss, epoch)\n",
    "        tb.add_scalar('Validation Accuracy', valid_correct/len_val, epoch)\n",
    "\n",
    "\n",
    "        # print training/validation statistics\n",
    "        # calculate average loss over an epoch\n",
    "        train_loss = train_loss/len_train\n",
    "        valid_loss = valid_loss/len_val\n",
    "        train_loop.write(f'\\t\\tAvg training loss: {train_loss:.6f}\\tAvg validation loss: {valid_loss:.6f}')\n",
    "\n",
    "\n",
    "        # save model if validation loss has decreased\n",
    "        if valid_loss <= valid_loss_min:\n",
    "            train_loop.write(f'\\t\\tvalid_loss decreased ({valid_loss_min:.6f} --> {valid_loss:.6f})  saving model...')\n",
    "            torch.save(resnet18.state_dict(), f'./model/lr3e-5{comment}.pth')\n",
    "            valid_loss_min = valid_loss\n",
    "\n",
    "\n",
    "        test_loss, test_correct = 0, 0  # will be used to track the running test loss and correct\n",
    "        ##################\n",
    "        # test the model #\n",
    "        ##################\n",
    "        for batch in test_loader:\n",
    "            images, labels = batch[0].to(device), batch[1].to(device)  # load the batch to available device\n",
    "            preds = resnet18(images)  # forward pass\n",
    "            loss = criterion(preds, labels)  # calculate the loss\n",
    "\n",
    "            test_loss += loss.item() * labels.size(0)  # update the running loss\n",
    "            test_correct += get_num_correct(preds, labels)  # update running num correct\n",
    "\n",
    "        # add test loss and test accuracy for the current epoch to tensorboard\n",
    "        tb.add_scalar('Test Loss', test_loss, epoch)\n",
    "        tb.add_scalar('Test Accuracy', test_correct/len_test, epoch)"
   ]
  }
 ]
}