notebooks/pytorch_IBA_per_sample.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyTorch: IBA (Per-Sample Bottleneck)\n",
    "\n",
    "This notebook shows how to apply the Per-Sample Bottleneck to pretrained ImageNet models. \n",
    "\n",
    "Ensure that `./imagenet` points to your copy of the ImageNet dataset. \n",
    "\n",
    "You might want to create a symlink:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! ln -s /path/to/your/imagenet/folder/ imagenet "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# to set you cuda device\n",
    "%env CUDA_VISIBLE_DEVICES=0\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import torch\n",
    "import torchvision.models \n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.datasets import ImageFolder\n",
    "from torchvision.transforms import Compose, CenterCrop, ToTensor, Resize, Normalize\n",
    "import matplotlib.pyplot as plt \n",
    "import os\n",
    "from tqdm import tqdm_notebook\n",
    "import json\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import sys\n",
    "\n",
    "try:\n",
    "    import IBA\n",
    "except ModuleNotFoundError:\n",
    "    sys.path.insert(0, '..')\n",
    "    import IBA\n",
    "    \n",
    "from IBA.pytorch import IBA, tensor_to_np_img\n",
    "from IBA.utils import plot_saliency_map"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading Data and Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imagenet_dir = './imagenet'\n",
    "\n",
    "dev = torch.device('cuda:0')\n",
    "\n",
    "# select a model to analyse\n",
    "# model = torchvision.models.vgg16(pretrained=True)\n",
    "model = torchvision.models.resnet50(pretrained=True)\n",
    "# model = torchvision.models.inception_v3(pretrained=True)\n",
    "model.to(dev).eval()\n",
    "\n",
    "# load the data\n",
    "if type(model) == torchvision.models.inception.Inception3:\n",
    "    image_size = 299\n",
    "else:\n",
    "    image_size = 224\n",
    "    \n",
    "valset = ImageFolder(\n",
    "    os.path.join(imagenet_dir, 'validation'),\n",
    "    transform=Compose([\n",
    "        CenterCrop(256), Resize(image_size), ToTensor(), \n",
    "        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "    ]))\n",
    "\n",
    "trainset = ImageFolder(\n",
    "    os.path.join(imagenet_dir, 'train'),\n",
    "    transform=Compose([\n",
    "        CenterCrop(256), Resize(image_size), ToTensor(), \n",
    "        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "    ]))\n",
    "\n",
    "with open('imagenet_class_index.json') as f:\n",
    "    idx2class = {int(k): v[1] for k, v in json.load(f).items()}\n",
    "\n",
    "trainloader = DataLoader(trainset, batch_size=64, shuffle=True, num_workers=4)\n",
    "img, target = valset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Insert the bottleneck into the model\n",
    "\n",
    "You can experiment with the location of the bottleneck."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if type(model) == torchvision.models.vgg.VGG:\n",
    "    iba = IBA(model.features[17])\n",
    "elif type(model) == torchvision.models.resnet.ResNet:\n",
    "    iba = IBA(model.layer2)\n",
    "elif type(model) == torchvision.models.inception.Inception3:\n",
    "    iba = IBA(model.Mixed_5b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iba.reset_estimate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Estimate Mean and Variance\n",
    "\n",
    "Here, we estimate the mean and variances of the feature map. It is important for measuring the amount of information transmitted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iba.estimate(model, trainloader, device=dev, n_samples=10000, progbar=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "neuron = (12, 3, 4)\n",
    "print(\"Neuron at position {:} has mean {:.2f} and std {:.2f}\".format(\n",
    "    neuron, iba.estimator.mean()[neuron],  iba.estimator.std()[neuron]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iba.estimator.n_samples()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Display Heatmaps for some random samples\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 5, figsize=(20, 6))\n",
    "np.random.seed(0)\n",
    "for ax, sample_idx in zip(axes.flatten(), np.random.choice(50000, 10)):\n",
    "    img, target = valset[sample_idx]\n",
    "    img = img[None].to(dev)\n",
    "    \n",
    "    # execute the model on a given sample and return the target NLL\n",
    "    model_loss_closure = lambda x: -torch.log_softmax(model(x), 1)[:, target].mean()\n",
    "    \n",
    "    # generate the heatmap\n",
    "    heatmap = iba.analyze(img, model_loss_closure)\n",
    "    \n",
    "    # reverse the data pre-processing for plotting the original image\n",
    "    np_img = tensor_to_np_img(img[0])\n",
    "    \n",
    "    # show the heatmap\n",
    "    plot_saliency_map(heatmap, np_img,  ax=ax)\n",
    "    ax.set_title(idx2class[target])\n",
    "    \n",
    "fig.suptitle(\"model: {}\".format(type(model).__name__))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Monkey image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = np.array(Image.open(\"./monkeys.jpg\"))\n",
    "img = (img.transpose(2, 0, 1) / 255)\n",
    "target = 382  # 382: squirrel monkey\n",
    "\n",
    "# preprocess image\n",
    "img  = Compose([\n",
    "    Resize(image_size), ToTensor(),  \n",
    "    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])(Image.open(\"./monkeys.jpg\"))\n",
    "\n",
    "model_loss_closure = lambda x: -torch.log_softmax(model(x), 1)[:, target].mean()\n",
    "heatmap = iba.analyze(img[None].to(dev), model_loss_closure) \n",
    "ax = plot_saliency_map(heatmap, tensor_to_np_img(img))\n",
    "_ = ax.set_title(idx2class[target])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}