calibration_temperature_scaling.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f1e54ede-f04a-4726-b23c-cb2efcf450a0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "\n",
    "import argparse\n",
    "import csv\n",
    "import os\n",
    "import sys\n",
    "import tifffile\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import glob\n",
    "import pickle\n",
    "\n",
    "import datetime\n",
    "import math\n",
    "import json\n",
    "import time\n",
    "\n",
    "import timm\n",
    "from timm.models import create_model\n",
    "from timm.scheduler import CosineLRScheduler\n",
    "from timm.optim import create_optimizer\n",
    "from timm.utils import NativeScaler, get_state_dict\n",
    "\n",
    "from torch.utils.data import Dataset, random_split, DataLoader\n",
    "from PIL import Image\n",
    "import torchvision.models as models\n",
    "import torchvision.transforms as T\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "from torchvision.utils import make_grid\n",
    "\n",
    "import torch.multiprocessing as mp\n",
    "import torch.distributed as dist\n",
    "from torch.nn.parallel import DistributedDataParallel as DDP\n",
    "from torch.utils.data.distributed import DistributedSampler\n",
    "\n",
    "from customloss import CustomLoss\n",
    "from datamgr import *\n",
    "import utils\n",
    "from engine import *\n",
    "from dataset import *\n",
    "from models import *\n",
    "from samplers import RASampler\n",
    "\n",
    "import copy\n",
    "from pathlib import Path\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "75d30986-450b-47d8-82bf-98508a06b5b0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def parse_args():\n",
    "    parser = argparse.ArgumentParser(description= 'CBNA training script')\n",
    "    parser.add_argument('--model_type'       , default='MLP', help='CNN/MLP/ViT/Fusion')\n",
    "    #parser.add_argument('--data_split'       , default='train', help='train/test')\n",
    "    #parser.add_argument('--get'       , default='predictions', help='features/predictions')\n",
    "    parser.add_argument('--data_set'     , default='CBNA')\n",
    "    parser.add_argument('--data_path'     , default='Data/irc_patches/')\n",
    "    parser.add_argument('--file_path'     , default='datafile/')\n",
    "    parser.add_argument('--output_dir', default='checkpoints/', help='path where to save, empty for no saving')\n",
    "    parser.add_argument('--device', default='cpu', help='device to use for training / testing')\n",
    "    parser.add_argument('--batch_size', default=128, type=int, help='Per-GPU batch-size : number of distinct images loaded on one GPU.')\n",
    "    parser.add_argument('--input_size'  , default=224, type=int, help ='Image size for training')\n",
    "    parser.add_argument('--drop_path'   , type=float, default=0.05)\n",
    "    parser.add_argument('--loss_fn'     , default='focal', help='bce/focal')\n",
    "    parser.add_argument('--weighted'    , type=utils.bool_flag, default=True)\n",
    "    parser.add_argument('--eval_crop_ratio', default=0.875, type=float, help=\"Crop ratio for evaluation\")\n",
    "    parser.add_argument('--num_workers', default=8, type=int, help='Number of data loading workers per GPU.')\n",
    "    parser.add_argument('--seed'       , default=0, type=int, help='seed')\n",
    "    parser.add_argument('--eval'        , action='store_true')\n",
    "    parser.add_argument('--thres_method'   , default='global', help='global/adaptive')\n",
    "    parser.add_argument('--threshold'      , default=0.5, type=float, help='compute TSS, 0.45 for CNN & Fusion, 0.5 for MLP')\n",
    "\n",
    "    return parser.parse_args('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4c157b65-5202-4092-a062-d440756ada26",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of labels:  2522\n"
     ]
    }
   ],
   "source": [
    "args = parse_args()\n",
    "output_dir = Path(args.output_dir).joinpath('{}/{}/'.format(args.data_set, args.model_type))\n",
    "args.output_dir = output_dir\n",
    "if not os.path.isdir(args.output_dir):\n",
    "    try:\n",
    "        os.makedirs(args.output_dir, exist_ok = True)\n",
    "        print(\"Directory '%s' created successfully\" %args.output_dir)\n",
    "    except OSError as error:\n",
    "        print(\"Directory '%s' can not be created\")\n",
    "\n",
    "if args.data_set == 'CBNA':\n",
    "    args.classes = list(np.genfromtxt(Path(args.file_path).joinpath('classes.txt'), dtype='str'))\n",
    "    args.num_classes = len(args.classes)\n",
    "    print('number of labels: ', args.num_classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9dafc29b-5ee9-40f5-8bee-dd02956d64fc",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data loaded: there are 98763 train images, 15980 test images.\n",
      "Creating model: MLP\n",
      "number of params: 401498\n",
      "best_file checkpoints/CBNA/MLP/best.pth\n",
      "model loaded\n"
     ]
    }
   ],
   "source": [
    "utils.fix_random_seeds(args.seed)\n",
    "cudnn.benchmark = True\n",
    "\n",
    "#print(args)\n",
    "\n",
    "num_tasks = utils.get_world_size()\n",
    "global_rank = utils.get_rank()\n",
    "\n",
    "meta_features = None\n",
    "if args.model_type==\"MLP\" or args.model_type==\"Fusion\":\n",
    "    meta_features = ['LANDOLT_MOIST',\n",
    "'N_prct', 'pH', 'CN', 'TMeanY', 'TSeason', 'PTotY', 'PSeason', 'RTotY',\n",
    "'RSeason', 'AMPL', 'LENGTH', 'eauvive', 'clay', 'silt', 'sand', 'cv_alti']\n",
    "\n",
    "if args.model_type==\"MLP\":\n",
    "    dataset_train = CBNA(args, split='train', meta_features=meta_features)\n",
    "    dataset_val = CBNA(args, split='val', meta_features=meta_features)\n",
    "    dataset_test = CBNA(args, split='test', meta_features=meta_features)\n",
    "else: \n",
    "    dataset_train = CBNA(args, split='train', meta_features=meta_features, transform=build_transform(args, aug=False))\n",
    "    dataset_val = CBNA(args, split='val', meta_features=meta_features, transform=build_transform(args, aug=False))\n",
    "    dataset_test = CBNA(args, split='test', meta_features=meta_features, transform=build_transform(args, aug=False))\n",
    "\n",
    "print(f\"Data loaded: there are {len(dataset_train)} train images, {len(dataset_test)} test images.\")\n",
    "\n",
    "if args.model_type==\"MLP\" or args.model_type==\"Fusion\":\n",
    "    args.n_meta_features = len(dataset_train.meta_features)\n",
    "\n",
    "args.train_label_cnt = dataset_train.get_label_count()\n",
    "args.test_label_cnt = dataset_test.get_label_count()\n",
    "\n",
    "train_sampler = torch.utils.data.SequentialSampler(dataset_train)\n",
    "train_loader = DataLoader(dataset_train, sampler=train_sampler, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=True, drop_last=False)\n",
    "val_sampler = torch.utils.data.SequentialSampler(dataset_val)\n",
    "val_loader = DataLoader(dataset_val, sampler=val_sampler, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=True, drop_last=False)\n",
    "test_sampler = torch.utils.data.SequentialSampler(dataset_test)\n",
    "test_loader = DataLoader(dataset_test, sampler=test_sampler, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=True, drop_last=False)\n",
    "\n",
    "print(f\"Creating model: {args.model_type}\")\n",
    "if args.model_type == 'MLP':\n",
    "    model = Mlp_CBNA(args, out_dim=args.num_classes)\n",
    "elif args.model_type == 'CNN':\n",
    "    model = Resnet_CBNA(args, out_dim=args.num_classes)\n",
    "elif args.model_type == 'Fusion':\n",
    "    model = Fusion_CBNA(args, out_dim=args.num_classes)\n",
    "elif args.model_type == 'ViT':\n",
    "    model = Vit_CBNA(args, out_dim=args.num_classes)\n",
    "model.to(args.device)\n",
    "\n",
    "n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "print('number of params:', n_parameters)\n",
    "\n",
    "reload_file = utils.get_best_file(args.output_dir)        \n",
    "print(\"best_file\" , reload_file)\n",
    "checkpoint = torch.load(reload_file, map_location=torch.device('cpu'))\n",
    "model.load_state_dict(checkpoint['model'])\n",
    "print(\"model loaded\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9a3785a0-8426-43dd-a3ad-f2bfb7d3651c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 125/125 [00:03<00:00, 31.99it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'eval_loss': 0.00020920961105730385,\n",
       " 'macro_tss': 0.6960844993591309,\n",
       " 'micro_tss': 0.714148998260498,\n",
       " 'weighted_tss': 0.6047736406326294}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "criterion = CustomLoss(args.loss_fn, args)\n",
    "evaluate(model, test_loader, criterion, args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "32135f5d-b03d-4813-b8c4-e5aa628f990b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ModelWithTemperature(nn.Module):\n",
    "    \"\"\"\n",
    "    A thin decorator, which wraps a model with temperature scaling\n",
    "    model (nn.Module):\n",
    "        A classification neural network\n",
    "        NB: Output of the neural network should be the classification logits,\n",
    "            NOT the softmax (or log softmax)!\n",
    "    \"\"\"\n",
    "    def __init__(self, model):\n",
    "        super(ModelWithTemperature, self).__init__()\n",
    "        self.model = model\n",
    "        self.temperature = nn.Parameter(torch.ones(len(args.classes))*0.5)\n",
    "\n",
    "    def forward(self, input):\n",
    "        logits = self.model(input)\n",
    "        return self.temperature_scale(logits)\n",
    "\n",
    "    def temperature_scale(self, logits):\n",
    "        \"\"\"\n",
    "        Perform temperature scaling on logits\n",
    "        \"\"\"\n",
    "        # Expand temperature to match the size of logits\n",
    "        temperature = self.temperature\n",
    "        return logits / temperature\n",
    "\n",
    "    # This function probably should live outside of this class, but whatever\n",
    "    def set_temperature(self, val_loader):\n",
    "        \n",
    "        criterion = CustomLoss(args.loss_fn, args)\n",
    "\n",
    "        # First: collect all the logits and labels for the validation set\n",
    "        logits_list = []\n",
    "        labels_list = []\n",
    "        with torch.no_grad():\n",
    "            for data, targets in val_loader:\n",
    "                if args.model_type==\"MLP\" or args.model_type==\"CNN\" or args.model_type==\"ViT\":\n",
    "                    samples = data\n",
    "                    samples = samples.to(args.device, non_blocking=True)\n",
    "                    logits = self.model(samples)\n",
    "                elif args.model_type==\"Fusion\":\n",
    "                    samples_img, samples_meta = data\n",
    "                    samples_img = samples_img.to(args.device, non_blocking=True)\n",
    "                    samples_meta = samples_meta.to(args.device, non_blocking=True)\n",
    "                    logits = self.model(samples_img, samples_meta)\n",
    "                \n",
    "                targets = targets.to(args.device, non_blocking=True)\n",
    "                logits_list.append(logits)\n",
    "                labels_list.append(targets)\n",
    "            logits = torch.cat(logits_list)\n",
    "            labels = torch.cat(labels_list)\n",
    "\n",
    "        loss_before_temperature = criterion(logits, labels).item()\n",
    "\n",
    "        # Next: optimize the temperature w.r.t\n",
    "        optimizer = torch.optim.AdamW([self.temperature], lr=1e-1, weight_decay=5e-2)\n",
    "\n",
    "        for i in range(15):\n",
    "            optimizer.zero_grad()\n",
    "            loss = criterion(self.temperature_scale(logits), labels)\n",
    "            print('loss', loss.item())\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        loss_after_temperature = criterion(self.temperature_scale(logits), labels).item()\n",
    "        \n",
    "        print('Before temperature scaling: ', loss_before_temperature)\n",
    "        #print('Optimal temperature: ', self.temperature)\n",
    "        print('After temperature scaling: ', loss_after_temperature)\n",
    "\n",
    "        return self"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9e0ff0fe-df9e-4626-90dd-af4dba14cde6",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss 0.0003092587285209447\n",
      "loss 0.0002684619394131005\n",
      "loss 0.00024264470266643912\n",
      "loss 0.00022566104598809034\n",
      "loss 0.0002141133154509589\n",
      "loss 0.00020603384473361075\n",
      "loss 0.00020023580873385072\n",
      "loss 0.0001959777728188783\n",
      "loss 0.00019278208492323756\n",
      "loss 0.00019033285207115114\n",
      "loss 0.00018841648125089705\n",
      "loss 0.00018688567797653377\n",
      "loss 0.00018563726916909218\n",
      "loss 0.00018459807324688882\n",
      "loss 0.00018371552869211882\n",
      "Before temperature scaling:  0.000200928290723823\n",
      "After temperature scaling:  0.00018295169866178185\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ModelWithTemperature(\n",
       "  (model): Mlp_CBNA(\n",
       "    (meta): Sequential(\n",
       "      (0): Linear(in_features=17, out_features=512, bias=True)\n",
       "      (1): BatchNorm1d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): ReLU(inplace=True)\n",
       "      (3): Dropout(p=0.3, inplace=False)\n",
       "      (4): Linear(in_features=512, out_features=128, bias=True)\n",
       "      (5): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (6): ReLU(inplace=True)\n",
       "    )\n",
       "    (fc): Linear(in_features=128, out_features=2522, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_w_tem = ModelWithTemperature(model).to(args.device)\n",
    "model_w_tem.set_temperature(val_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d00ad0ce-afa1-4e53-9d48-f64699f32a1a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 125/125 [00:03<00:00, 31.74it/s]\n",
      "TSS after temperature scaling:  {'eval_loss': 0.000195201369933784, 'macro_tss': 0.6960844993591309, 'micro_tss': 0.714148998260498, 'weighted_tss': 0.6047736406326294}\n"
     ]
    }
   ],
   "source": [
    "TP, TN, FP, FN = 0., 0., 0., 0.\n",
    "label_cnt = 0.\n",
    "eval_loss = []\n",
    "model_w_tem.eval()\n",
    "with torch.no_grad():\n",
    "    for data, targets in tqdm(test_loader, file=sys.stdout):\n",
    "        if args.model_type==\"MLP\" or args.model_type==\"CNN\" or args.model_type==\"ViT\":\n",
    "            samples = data\n",
    "            samples = samples.to(args.device, non_blocking=True)\n",
    "            inputs = samples\n",
    "            outputs = model_w_tem(inputs)\n",
    "        elif args.model_type==\"Fusion\":\n",
    "            samples_img, samples_meta = data\n",
    "            samples_img = samples_img.to(args.device, non_blocking=True)\n",
    "            samples_meta = samples_meta.to(args.device, non_blocking=True)\n",
    "            inputs = (samples_img, samples_meta)\n",
    "            outputs = model_w_tem(inputs)\n",
    "\n",
    "        targets = targets.to(args.device, non_blocking=True)\n",
    "        label_cnt += targets.sum(0).float()\n",
    "        \n",
    "        loss = criterion(outputs, targets)\n",
    "        eval_loss.append(loss.item())\n",
    "\n",
    "        if args.thres_method == 'adaptive':\n",
    "            thresholds = torch.load(args.output_dir.joinpath('thresholds_train.pth'))\n",
    "            thresholds = torch.tensor(thresholds).to(args.device)\n",
    "        elif args.thres_method == 'global':\n",
    "            thresholds = args.threshold\n",
    "\n",
    "        tp, tn, fp, fn = compute_scores(torch.sigmoid(outputs), targets, thresholds)\n",
    "        TP += tp\n",
    "        TN += tn\n",
    "        FP += fp\n",
    "        FN += fn\n",
    "\n",
    "    eval_loss = torch.tensor(eval_loss).mean().item()\n",
    "    weight = label_cnt / label_cnt.sum()\n",
    "    macro_tss, micro_tss, weighted_tss = compute_metrics(TP, TN, FP, FN, weight=weight)\n",
    "\n",
    "    eval_stats = {'eval_loss': eval_loss, 'macro_tss': macro_tss, 'micro_tss': micro_tss, 'weighted_tss': weighted_tss}\n",
    "    print(\"TSS after temperature scaling: \", eval_stats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "607c48e0-23b2-4904-8d07-09ad89ee4cbf",
   "metadata": {},
   "outputs": [],
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}