ImageNet_ResNet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3c807aa-554f-471a-a46c-22b138e6849f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# models and data processing from https://github.com/SamsungLabs/pytorch-ensembles/blob/master\n",
    "\n",
    "import os\n",
    "import urllib.request\n",
    "from types import SimpleNamespace\n",
    "from urllib.error import HTTPError\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pytorch_lightning as pl\n",
    "import seaborn as sns\n",
    "import tabulate\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as data\n",
    "import torchvision\n",
    "\n",
    "%matplotlib inline\n",
    "from IPython.display import HTML, display, set_matplotlib_formats\n",
    "from PIL import Image\n",
    "from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint\n",
    "from torchvision import transforms\n",
    "from torchvision.datasets import CIFAR10\n",
    "\n",
    "import torch.nn.functional as F\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "from utils import weighted_mean, weighted_sd, read_results\n",
    "from torch_utils import d_LSE, BI_LSE, BI_thresh_acc_DE, BI_thresh_acc_TTA, Conf_thresh_acc\n",
    "\n",
    "set_matplotlib_formats(\"svg\", \"pdf\")  # For export\n",
    "matplotlib.rcParams[\"lines.linewidth\"] = 2.0\n",
    "sns.reset_orig()\n",
    "\n",
    "os.makedirs('plots', exist_ok=True)\n",
    "ds_name = 'ImageNet-C'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4999996f-4e4c-46ee-a670-fbcd81372cc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "prefix = '../saved_models/deepens_imagenet/ImageNet-ResNet50-{}-1.pth.tar'\n",
    "suffix_list = [\n",
    "    '052e7f78e4db--1564492444',\n",
    "    '1132c260ef75--1564493784',\n",
    "    '2f817072e8da--1564493734',\n",
    "    '3177c697fbf4--1564495013',\n",
    "    '628e11f9fd67--1564481099',\n",
    "    '743e10f26a38--1564493675',\n",
    "    '7ded66ec9900--1564481097',\n",
    "    '8fc5076a66c9--1564481079',\n",
    "    'a58ab8dd26fc--1564492521',\n",
    "    'a80e40d84db2--1564492573',\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "642d5bc2-f3ee-4696-baaa-84478cc928fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "models = []\n",
    "for suffix in suffix_list:\n",
    "    checkpoint = torch.load(prefix.format(suffix))\n",
    "    model = torchvision.models.resnet50()\n",
    "    model = torch.nn.DataParallel(model, device_ids = [0, 1, 2, 3]).cuda()\n",
    "    model.load_state_dict(checkpoint['state_dict'])\n",
    "    models.append(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df25a958-7713-4b70-ba0e-22902ade45e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_PATH = os.environ.get(\"PATH_DATASETS\", \"../data/\")\n",
    "DATA_MEANS = (0.485, 0.456, 0.406)\n",
    "DATA_STD = (0.229, 0.224, 0.225)\n",
    "batch_size = 64\n",
    "\n",
    "test_transform = transforms.Compose([\n",
    "    transforms.Resize(256),\n",
    "    transforms.CenterCrop(224),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(DATA_MEANS, DATA_STD)]\n",
    ")\n",
    "\n",
    "val_dataset = torchvision.datasets.ImageNet(root=DATASET_PATH, split='val', transform=test_transform)\n",
    "print(len(val_dataset))\n",
    "pl.seed_everything(42)\n",
    "val_set, test_set = torch.utils.data.random_split(val_dataset, [10000, 40000])\n",
    "\n",
    "# We define a set of data loaders that we can use for various purposes later.\n",
    "val_loader = data.DataLoader(val_set, batch_size=batch_size, shuffle=False, drop_last=False, num_workers=4, pin_memory=True)\n",
    "test_loader = data.DataLoader(test_set, batch_size=batch_size, shuffle=False, drop_last=False, num_workers=4, pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2aeeda16-a2bd-4bf1-8a3d-9897428148b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ImageNet-C\n",
    "\n",
    "def imshow(loader):\n",
    "    dataiter = iter(loader)\n",
    "    images, labels = dataiter.next()\n",
    "    img = torchvision.utils.make_grid(images)\n",
    "    img = img / 2 + 0.5     # unnormalize\n",
    "    npimg = img.numpy()\n",
    "    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
    "    plt.show()\n",
    "\n",
    "def get_loader(corruption, severity = 5):\n",
    "    dataset = torchvision.datasets.ImageFolder(\n",
    "        root='../data/ImageNet-C/' + corruption + '/' + str(severity),\n",
    "        transform=test_transform\n",
    "    )\n",
    "    # to avoid data leakage, we require an identical val / test set split as above\n",
    "    pl.seed_everything(42)\n",
    "    _, test_set = torch.utils.data.random_split(dataset, [10000, 40000])\n",
    "\n",
    "\n",
    "    return torch.utils.data.DataLoader(\n",
    "        test_set, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fa37d89-5ab5-47cf-ba2d-ed106c4ac9f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# checking for data leak\n",
    "#imshow(test_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9718a64b-1849-47c0-a3eb-4c1e35383d33",
   "metadata": {},
   "outputs": [],
   "source": [
    "# No data leak observable\n",
    "#imshow(get_loader('brightness'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cb11952-92ab-4964-8f53-ead2ef5b23b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# this cell takes >10 hours (single RTX 5000)\n",
    "qs = np.arange(10, -1, -1) / 10\n",
    "corruptions = [\n",
    "    'None',\n",
    "    'brightness',\n",
    "    'fog',\n",
    "    'glass_blur',\n",
    "    'pixelate',\n",
    "    'spatter',\n",
    "    'contrast',\n",
    "    'frost',\n",
    "    'impulse_noise',\n",
    "    'saturate',\n",
    "    'speckle_noise',\n",
    "    'defocus_blur',\n",
    "    'gaussian_blur',\n",
    "    'jpeg_compression',\n",
    "    'shot_noise',\n",
    "    'zoom_blur',\n",
    "    'elastic_transform',\n",
    "    'gaussian_noise',\n",
    "    'motion_blur',\n",
    "    'snow',\n",
    "]\n",
    "target = 'Accuracy'\n",
    "\n",
    "for type_ in ['DE', 'TTA', 'Conf']:\n",
    "    if type_=='DE':\n",
    "        experiment = BI_thresh_acc_DE(models, val_loader, qs=qs, device='cuda:0')\n",
    "    elif type_=='TTA':\n",
    "        experiment = BI_thresh_acc_TTA(models[0], val_loader, qs=qs, n_ens=20, device='cuda:0')\n",
    "    elif type_=='Conf':\n",
    "        experiment = Conf_thresh_acc(models[0], qs=qs, device='cuda:0')\n",
    "\n",
    "    for severity in tqdm(range(1, 6)):\n",
    "        loaders = []\n",
    "        for corr in tqdm(corruptions, leave=False):\n",
    "            if corr == 'None':\n",
    "                loaders += [test_loader]\n",
    "            else:\n",
    "                loaders += [get_loader(corr, severity=severity)]\n",
    "\n",
    "        results_BI_ood = experiment.run(loaders)\n",
    "        results_BI_ood = results_BI_ood.reset_index(drop=True)\n",
    "        os.makedirs('results/{}/{}/'.format(ds_name, target), exist_ok=True)\n",
    "        results_BI_ood.to_pickle('results/{}/{}/all/{}_sev{}.pkl'.format(ds_name, target, type_, severity))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23d6baa5-c2ac-4da1-acab-5513e8d26b44",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Figure 7\n",
    "sns.set_palette(['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf'])\n",
    "target = 'Accuracy'\n",
    "hue = 'Corruption Severity'\n",
    "\n",
    "unc_type = 'Conf'\n",
    "results_Conf_all = []\n",
    "for severity in range(1, 6):\n",
    "    results_Conf_all += [read_results(unc_type=unc_type, ds_name=ds_name, severity=severity, target=target)]\n",
    "results_Conf_all = pd.concat(results_Conf_all)\n",
    "\n",
    "unc_type = 'DE'\n",
    "results_DE_all = []\n",
    "for severity in range(1, 6):\n",
    "    results_DE_all += [read_results(unc_type=unc_type, ds_name=ds_name, severity=severity, target=target)]\n",
    "results_DE_all = pd.concat(results_DE_all)\n",
    "\n",
    "unc_type = 'TTA'\n",
    "results_TTA_all = []\n",
    "for severity in range(1, 6):\n",
    "    results_TTA_all += [read_results(unc_type=unc_type, ds_name=ds_name, severity=severity, target=target)]\n",
    "results_TTA_all = pd.concat(results_TTA_all)\n",
    "\n",
    "\n",
    "plt.rcParams.update({'font.size': 14, 'figure.dpi': 1000})\n",
    "figure, axs = plt.subplots(1, 3, figsize=(15, 4))\n",
    "\n",
    "hue = 'Corruption Severity'\n",
    "\n",
    "ax = axs[0]\n",
    "plot = sns.lineplot(\n",
    "    data=results_Conf_all, x='Validation set quantile', y='Accuracy', hue=hue, ax=ax, linewidth=3,\n",
    "    estimator=weighted_mean, errorbar=weighted_sd, legend=False)\n",
    "#plot.axhline(1., c='black', ls='dashed')\n",
    "plot.invert_xaxis()\n",
    "ax.set_ylim(0.4, 1.0)\n",
    "ax.set_xlim(0.9, 0.4)\n",
    "ax.set_title('Confidence Score')\n",
    "sns.despine()\n",
    "\n",
    "ax = axs[1]\n",
    "plot = sns.lineplot(\n",
    "    data=results_TTA_all, x='Validation set quantile', y='Accuracy', hue=hue, ax=ax, linewidth=3, \n",
    "    estimator=weighted_mean, errorbar=weighted_sd, legend=False)\n",
    "plot.invert_xaxis()\n",
    "#plot.axhline(1., c='black', ls='dashed')\n",
    "ax.set_ylim(0.4, 1.0)\n",
    "ax.set_xlim(0.9, 0.4)\n",
    "ax.set_ylabel('')\n",
    "ax.set_title('Bregman Information (TTA)')\n",
    "sns.despine()\n",
    "\n",
    "ax = axs[2]\n",
    "plot = sns.lineplot(\n",
    "    data=results_DE_all, x='Validation set quantile', y='Accuracy', hue=hue, ax=ax, linewidth=3, \n",
    "    estimator=weighted_mean, errorbar=weighted_sd)\n",
    "plot.invert_xaxis()\n",
    "#plot.axhline(1., c='black', ls='dashed')\n",
    "ax.set_ylim(0.4, 1.0)\n",
    "ax.set_xlim(0.9, 0.4)\n",
    "ax.set_ylabel('')\n",
    "ax.set_title('Bregman Information (DE)')\n",
    "sns.despine()\n",
    "\n",
    "ax.legend(bbox_to_anchor=(1.02, 0.15), loc='upper left', borderaxespad=0)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig('plots/resnet_imgnet-C_all.png', bbox_inches='tight')\n",
    "plt.show();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3884f924-cf42-46f3-a856-1b96fa628398",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}