EffNet_Classification.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "UNet_Classification.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qPH-VDcSvhh0"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "import sklearn\n",
        "from skimage.io import imread \n",
        "from PIL import Image\n",
        "import csv\n",
        "import torch\n",
        "import copy\n",
        "import torchvision\n",
        "from torchvision.datasets import ImageFolder\n",
        "from torch.utils.data import DataLoader, random_split\n",
        "from torch.utils import data\n",
        "import torchvision.transforms as transf\n",
        "import torch.nn.functional as F\n",
        "from torchvision.transforms.transforms import RandomAdjustSharpness\n",
        "import torch.nn as nn\n",
        "import random\n",
        "import cv2\n",
        "import math\n",
        "import os\n",
        "from torch import optim\n",
        "from torch.nn.modules.conv import ConvTranspose2d\n",
        "from torch.autograd import Function\n",
        "import glob"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v5E9wxdn0s--"
      },
      "outputs": [],
      "source": [
        "# Hyperparameter Tuning\n",
        "epoch_n = 50\n",
        "l_rate = 0.0001\n",
        "batch_size_tr = 15\n",
        "batch_size_val = 15"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Structure of dataset\n",
        "\n",
        "--- /train\n",
        "------/class_1\n",
        "---------/mask_gen1.png\n",
        "---------/mask_gen2.png\n",
        "\n",
        "------/class_2\n",
        "---------/mask_gen3.png\n",
        "---------/mask_gen4.png\n",
        "\n",
        "------/class_3\n",
        "...\n",
        "\n",
        "--- /val\n",
        "------/class_1\n",
        "---------/mask_gen6.png\n",
        "---------/mask_gen7.png\n",
        "\n",
        "------/class_2\n",
        "---------/mask_gen8.png\n",
        "---------/mask_gen9.png\n",
        "\n",
        "------/class_3\n",
        "..."
      ],
      "metadata": {
        "id": "rIU_LrOjwwEE"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6xczE6sl0s--"
      },
      "outputs": [],
      "source": [
        "# Augmentation of dataset\n",
        "\n",
        "train_transform = transf.Compose([\n",
        "    transf.Resize((380,380)),\n",
        "    transf.RandomHorizontalFlip(p=0.5),\n",
        "    transf.RandomRotation(degrees=45),\n",
        "    transf.RandomVerticalFlip(p=0.5),\n",
        "    transf.ColorJitter(brightness=0.5),\n",
        "    transf.RandomAdjustSharpness(sharpness_factor=5),\n",
        "    transf.RandomAutocontrast(),\n",
        "    transf.ToTensor()\n",
        "])\n",
        "\n",
        "val_transform = transf.Compose([\n",
        "    transf.Resize((224,224)),\n",
        "    transf.ToTensor()\n",
        "])\n",
        "train_ds = ImageFolder('/content/drive/MyDrive/PH2/train_classifier/', transform=train_transform)\n",
        "val_ds = ImageFolder('/content/drive/MyDrive/PH2/val_classifier/', transform=val_transform)\n",
        "#test_ds = ImageFolder('/content/drive/MyDrive/Lung_Carcinoma/data_folder/test', transform=val_transform)\n",
        "train_load = DataLoader(dataset=train_ds, batch_size=batch_size_tr, shuffle=True, num_workers=2, drop_last=False)\n",
        "val_load = DataLoader(dataset=val_ds, batch_size=batch_size_val, shuffle=True, num_workers=2, drop_last=False)\n",
        "#test_load = DataLoader(dataset=test_ds, batch_size=batch_size_val, shuffle=False, drop_last=False)\n",
        "if torch.cuda.is_available():\n",
        "  device='cuda'\n",
        "else:\n",
        "  device='cpu'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Yk6otsaF4MJt"
      },
      "outputs": [],
      "source": [
        "# Deep Feature Extraction modifications\n",
        "\n",
        "class refined(nn.Module):\n",
        "  def __init__(self, out_classes=3):\n",
        "    super(refined,self).__init__()\n",
        "    eff_model = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_efficientnet_b4', pretrained=True)\n",
        "    self.initial = nn.Sequential(*list(eff_model.children())[:1])\n",
        "    self.ref_model_freeze = nn.Sequential(*list(eff_model.children())[1][:6])\n",
        "    self.ref_model = nn.Sequential(*list(eff_model.children())[1][6],\n",
        "                                   *list(eff_model.children())[2:][0])\n",
        "    self.pool = nn.AdaptiveAvgPool2d(1)\n",
        "    self.flat = nn.Flatten()\n",
        "    # self.pool = nn.AdaptiveAvgPool2d(output_size=1)\n",
        "    # self.flat = nn.Flatten()\n",
        "    self.fcc = nn.Sequential(\n",
        "                 nn.Linear(1792,out_classes))\n",
        "  \n",
        "  def forward(self,images):\n",
        "    raw_featr_freeze = self.ref_model_freeze(self.initial(images))\n",
        "    raw_featr = self.flat(self.pool(self.ref_model(raw_featr_freeze)))\n",
        "    output = self.fcc(raw_featr)\n",
        "    return raw_featr,output"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E5yTqJ4SbiZ5"
      },
      "outputs": [],
      "source": [
        "##### training #####\n",
        "\n",
        "def train(model, criterion_tr, criterion_val, optim, epoch_n):\n",
        "  best_loss=5 # ekhane previous best acc ta dibi in the format 0.XX\n",
        "  train_loss_list = []\n",
        "  val_loss_list = []\n",
        "  best_featr_tr = []\n",
        "  best_ver_labels_ext_tr = []\n",
        "  best_featr_val = []\n",
        "  best_ver_labels_ext_val = []\n",
        "  #best_model_wts = copy.deepcopy(model.state_dict())\n",
        "  for epoch in range(epoch_n):\n",
        "    featr_tensor_tr = np.zeros((1,1792))\n",
        "    labels_ext_tr = []\n",
        "    model.train()\n",
        "    running_train_loss = 0.0\n",
        "    running_train_acc = 0.0\n",
        "    for images,labels in train_load:\n",
        "      images = images.to(device)\n",
        "      labels = labels.to(device)\n",
        "      with torch.set_grad_enabled(True):\n",
        "        featr,output = model(images)\n",
        "        featr_tensor_tr = np.append(featr_tensor_tr,featr.cpu().detach().numpy(), axis=0)\n",
        "        labels_ext_tr = np.append(labels_ext_tr,labels.cpu().detach().numpy(),axis=0)\n",
        "        _,pred = torch.max(output,1)\n",
        "        loss = criterion_tr(output,labels,featr)\n",
        "        loss.backward()\n",
        "        optim.step()\n",
        "      optim.zero_grad()\n",
        "      running_train_loss += loss.item()*batch_size_tr\n",
        "      running_train_acc += torch.sum(pred==labels)\n",
        "    running_val_loss, running_val_acc, featr_tensor_val, labels_ext_val = eval(model, criterion_val)\n",
        "    epoch_train_loss = running_train_loss/len(train_ds)\n",
        "    epoch_train_acc = running_train_acc.double()/len(train_ds)\n",
        "    print(\"Epoch: {}\".format(epoch+1))\n",
        "    print('-'*10)\n",
        "    print('Train Loss: {:.4f}   Train Acc: {:.4f}'.format(epoch_train_loss,epoch_train_acc))\n",
        "    epoch_val_loss = running_val_loss/len(val_ds)\n",
        "    epoch_val_acc = running_val_acc.double()/len(val_ds)\n",
        "    print('Val Loss: {:.4f}   Val Acc: {:.4f}'.format(epoch_val_loss,epoch_val_acc))\n",
        "    print('\\n')\n",
        "    if best_loss > epoch_val_loss:\n",
        "      best_loss = epoch_val_loss\n",
        "      best_model_wts = copy.deepcopy(model.state_dict())\n",
        "      state={\n",
        "          \"model_state\":model.state_dict(),\n",
        "          \"optimizer_state\":optim.state_dict(),\n",
        "            }\n",
        "      torch.save(state,'/content/drive/MyDrive/PH2/model_wts/model_wts_effnet_PH2_run1.pth') # ekhane notun weight er src ta korbi \n",
        "      best_featr_tr = featr_tensor_tr\n",
        "      best_ver_labels_ext_tr = labels_ext_tr\n",
        "      best_featr_val = featr_tensor_val\n",
        "      best_ver_labels_ext_val = labels_ext_val\n",
        "\n",
        "    last_model_wts = copy.deepcopy(model.state_dict())\n",
        "    state_1={\n",
        "        \"model_state\":model.state_dict(),\n",
        "        \"optimizer_state\":optim.state_dict(),\n",
        "          }\n",
        "    torch.save(state_1,'/content/drive/MyDrive/PH2/model_wts/model_wts_effnet_PH2_run1_last.pth')\n",
        "    train_loss_list = train_loss_list + [epoch_train_loss]\n",
        "    val_loss_list = val_loss_list + [epoch_val_loss]\n",
        "  model = model.load_state_dict(best_model_wts)\n",
        "  print(\"The model with the best performance has a loss of :{:.4f}\".format(best_loss))\n",
        "  return model, best_featr_tr, best_ver_labels_ext_tr, best_featr_val, best_ver_labels_ext_val, train_loss_list, val_loss_list\n",
        "\n",
        "\n",
        "def eval(model, criterion_val):\n",
        "  model.eval()\n",
        "  featr_tensor_val = np.zeros((1,1792))\n",
        "  labels_ext_val = []\n",
        "  running_val_loss = 0.0\n",
        "  running_val_acc = 0.0\n",
        "  for images,labels in val_load:\n",
        "    images = images.to(device)\n",
        "    labels = labels.to(device)\n",
        "    featr,output = model(images)\n",
        "    featr_tensor_val = np.append(featr_tensor_val,featr.cpu().detach().numpy(), axis =0)\n",
        "    labels_ext_val = np.append(labels_ext_val,labels.cpu().detach().numpy(),axis=0)\n",
        "    _,pred = torch.max(output,1)\n",
        "    loss = criterion_val(output,labels,featr)\n",
        "    running_val_loss += loss.item()*batch_size_val\n",
        "    running_val_acc += torch.sum(pred==labels)\n",
        "  return running_val_loss, running_val_acc, featr_tensor_val, labels_ext_val"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qvpkyse8Dl3K"
      },
      "outputs": [],
      "source": [
        "# samples in each class for train and val\n",
        "\n",
        "cls_num_list_train=[57, 63, 32]\n",
        "cls_num_list_val=[23, 17, 8]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "87DRuNfQbHtc"
      },
      "outputs": [],
      "source": [
        "# Defining per class weights for CBReweight for train\n",
        "\n",
        "train_sampler = None\n",
        "beta = 0.9999\n",
        "effective_num = 1.0 - np.power(beta, cls_num_list_train)\n",
        "per_cls_weights = (1.0 - beta) / np.array(effective_num)\n",
        "per_cls_weights = per_cls_weights / np.sum(per_cls_weights) * len(cls_num_list_train)\n",
        "per_cls_weights_tr = torch.FloatTensor(per_cls_weights).to(device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9jabP4sKSim6"
      },
      "outputs": [],
      "source": [
        "# Defining per class weights for CBReweight for validation\n",
        "\n",
        "train_sampler = None\n",
        "beta = 0.9999\n",
        "effective_num = 1.0 - np.power(beta, cls_num_list_val)\n",
        "per_cls_weights = (1.0 - beta) / np.array(effective_num)\n",
        "per_cls_weights = per_cls_weights / np.sum(per_cls_weights) * len(cls_num_list_val)\n",
        "per_cls_weights_val = torch.FloatTensor(per_cls_weights).to(device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_iN-uziZJz9f"
      },
      "outputs": [],
      "source": [
        "num_classes = 3\n",
        "# IB_Focal Loss \n",
        "def ib_focal_loss(input_values, ib, gamma):\n",
        "    \"\"\"Computes the ib focal loss\"\"\"\n",
        "    p = torch.exp(-input_values)\n",
        "    loss = (1 - p) ** gamma * input_values * ib\n",
        "    return loss.mean()\n",
        "\n",
        "class IB_FocalLoss(nn.Module):\n",
        "    def __init__(self, weight, alpha=1000000, gamma=0):\n",
        "        super(IB_FocalLoss, self).__init__()\n",
        "        assert alpha > 0\n",
        "        self.alpha = alpha\n",
        "        self.epsilon = 0.001\n",
        "        self.weight = weight\n",
        "        self.gamma = gamma\n",
        "\n",
        "    def forward(self, input, target, features):\n",
        "        grads = torch.sum(torch.abs(F.softmax(input, dim=1) - F.one_hot(target, num_classes)),1) # N * 1\n",
        "        # print(grads.shape)\n",
        "        ib = grads*(torch.sum(features.reshape(-1)))\n",
        "        ib = self.alpha / (ib + self.epsilon)\n",
        "        return ib_focal_loss(F.cross_entropy(input, target, reduction='none', weight=self.weight), ib, self.gamma)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eI8fpaEIbHte"
      },
      "outputs": [],
      "source": [
        "model = refined()\n",
        "# state=torch.load('/content/drive/MyDrive/PH2/model_wts/model_wts_effnet_PH2_run1_last.pth') # ekhane previous weight ta src korbi \n",
        "# model.load_state_dict(state['model_state'])\n",
        "model = model.to(device)\n",
        "criterion_tr = IB_FocalLoss(per_cls_weights_tr)\n",
        "criterion_tr=criterion_tr.to(device)\n",
        "criterion_val = IB_FocalLoss(per_cls_weights_val)\n",
        "criterion_val=criterion_val.to(device)\n",
        "optim = torch.optim.Adam(model.parameters(), l_rate)\n",
        "# optim.load_state_dict(state['optimizer_state'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Fl8BCPTJ0s_A"
      },
      "outputs": [],
      "source": [
        "model, best_featr_tr, labels_tr, best_featr_val, labels_val, train_loss_list, val_loss_list = train(model, criterion_tr, criterion_val, optim, epoch_n)"
      ]
    }
  ]
}