featuremap_variance.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "651e668a",
   "metadata": {},
   "source": [
    "# Feature Map Variance"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce9d3b41",
   "metadata": {},
   "source": [
    "We measure feature map variance to demonstrate that ***MSAs aggregate feature maps and Convs (and MLPs) diversify them*** as shown in Fig 9 and Fig D.1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d1410e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "# check whether run in Colab\n",
    "root = \".\"\n",
    "if \"google.colab\" in sys.modules:\n",
    "    print(\"Running in Colab.\")\n",
    "    !pip3 install matplotlib\n",
    "    !pip3 install einops==0.4.1\n",
    "    !pip3 install timm==0.5.4\n",
    "    !git clone https://github.com/xxxnell/how-do-vits-work.git\n",
    "    root = \"./how-do-vits-work\"\n",
    "    sys.path.append(root)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6440ed60",
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "import copy\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "import models\n",
    "import ops.tests as tests\n",
    "import ops.datasets as datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3aa9899e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# config_path = \"%s/configs/cifar10_vit.yaml\" % root\n",
    "config_path = \"%s/configs/cifar100_vit.yaml\" % root\n",
    "# config_path = \"%s/configs/imagenet_vit.yaml\" % root\n",
    "\n",
    "with open(config_path) as f:\n",
    "    args = yaml.load(f)\n",
    "    print(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9087d0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_args = copy.deepcopy(args).get(\"dataset\")\n",
    "train_args = copy.deepcopy(args).get(\"train\")\n",
    "val_args = copy.deepcopy(args).get(\"val\")\n",
    "model_args = copy.deepcopy(args).get(\"model\")\n",
    "optim_args = copy.deepcopy(args).get(\"optim\")\n",
    "env_args = copy.deepcopy(args).get(\"env\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14844620",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dataset_train, dataset_test = datasets.get_dataset(**dataset_args, download=True)\n",
    "dataset_name = dataset_args[\"name\"]\n",
    "num_classes = len(dataset_train.classes)\n",
    "\n",
    "dataset_train = DataLoader(\n",
    "    dataset_train, \n",
    "    shuffle=True, \n",
    "    num_workers=train_args.get(\"num_workers\", 4), \n",
    "    batch_size=train_args.get(\"batch_size\", 128),\n",
    ")\n",
    "dataset_test = DataLoader(\n",
    "    dataset_test, \n",
    "    num_workers=val_args.get(\"num_workers\", 4), \n",
    "    batch_size=val_args.get(\"batch_size\", 32),\n",
    ")\n",
    "\n",
    "print(\"Train: %s, Test: %s, Classes: %s\" % (\n",
    "    len(dataset_train.dataset), \n",
    "    len(dataset_test.dataset), \n",
    "    num_classes\n",
    "))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad0723a5",
   "metadata": {},
   "source": [
    "## Load and Prepare Pretrained Models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a9da41c",
   "metadata": {},
   "source": [
    "Divide the pretrained models into the sequences of blocks (the cells below provide the snippets for ResNet-50 and ViT-Ti):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41dd592e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# download and load a pretrained model for CIFAR-100\n",
    "url = \"https://github.com/xxxnell/how-do-vits-work-storage/releases/download/v0.1/resnet_50_cifar100_691cc9a9e4.pth.tar\"\n",
    "path = \"checkpoints/resnet_50_cifar100_691cc9a9e4.pth.tar\"\n",
    "models.download(url=url, path=path)\n",
    "\n",
    "name = \"resnet_50\"\n",
    "model = models.get_model(name, num_classes=num_classes,  # timm does not provide a ResNet for CIFAR\n",
    "                         stem=model_args.get(\"stem\", False))\n",
    "map_location = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "checkpoint = torch.load(path, map_location=map_location)\n",
    "model.load_state_dict(checkpoint[\"state_dict\"])\n",
    "\n",
    "# model → blocks. `blocks` is a sequence of blocks\n",
    "blocks = [\n",
    "    model.layer0,\n",
    "    *model.layer1,\n",
    "    *model.layer2,\n",
    "    *model.layer3,\n",
    "    *model.layer4,\n",
    "    model.classifier,\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dece74f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "import timm\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "# download and load a pretrained model for CIFAR-100\n",
    "url = \"https://github.com/xxxnell/how-do-vits-work-storage/releases/download/v0.1/vit_ti_cifar100_9857b21357.pth.tar\"\n",
    "path = \"checkpoints/vit_ti_cifar100_9857b21357.pth.tar\"\n",
    "models.download(url=url, path=path)\n",
    "\n",
    "model = timm.models.vision_transformer.VisionTransformer(\n",
    "    num_classes=100, img_size=32, patch_size=2,  # for CIFAR\n",
    "    embed_dim=192, depth=12, num_heads=3, qkv_bias=False,  # for ViT-Ti \n",
    ")\n",
    "model.name = \"vit_ti\"\n",
    "models.stats(model)\n",
    "map_location = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "checkpoint = torch.load(path, map_location=map_location)\n",
    "model.load_state_dict(checkpoint[\"state_dict\"])\n",
    "\n",
    "\n",
    "# model → blocks. `blocks` is a sequence of blocks\n",
    "\n",
    "class PatchEmbed(nn.Module):\n",
    "    def __init__(self, model):\n",
    "        super().__init__()\n",
    "        self.model = copy.deepcopy(model)\n",
    "        \n",
    "    def forward(self, x, **kwargs):\n",
    "        x = self.model.patch_embed(x)\n",
    "        cls_token = self.model.cls_token.expand(x.shape[0], -1, -1)\n",
    "        x = torch.cat((cls_token, x), dim=1)\n",
    "        x = self.model.pos_drop(x + self.model.pos_embed)\n",
    "        return x\n",
    "    \n",
    "class Residual(nn.Module):\n",
    "    def __init__(self, *fn):\n",
    "        super().__init__()\n",
    "        self.fn = nn.Sequential(*fn)\n",
    "        \n",
    "    def forward(self, x, **kwargs):\n",
    "        return self.fn(x, **kwargs) + x\n",
    "    \n",
    "    \n",
    "class Lambda(nn.Module):\n",
    "    def __init__(self, fn):\n",
    "        super().__init__()\n",
    "        self.fn = fn\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.fn(x)\n",
    "\n",
    "\n",
    "def flatten(xs_list):\n",
    "    return [x for xs in xs_list for x in xs]\n",
    "\n",
    "\n",
    "# model → blocks. `blocks` is a sequence of blocks\n",
    "blocks = [\n",
    "    PatchEmbed(model),\n",
    "    *flatten([[Residual(b.norm1, b.attn), Residual(b.norm2, b.mlp)] \n",
    "              for b in model.blocks]),\n",
    "    nn.Sequential(model.norm, Lambda(lambda x: x[:, 0]), model.head),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fed203fb",
   "metadata": {},
   "source": [
    "## Plot Feature Map Variances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e32b0435",
   "metadata": {},
   "outputs": [],
   "source": [
    "xs, ys = next(iter(dataset_test))  # only one batch is used for simplicity\n",
    "\n",
    "# accumulate `latents` by collecting hidden states of a model\n",
    "latents = []\n",
    "with torch.no_grad():\n",
    "    for block in blocks:\n",
    "        xs = block(xs)\n",
    "        latents.append(xs)\n",
    "        \n",
    "if model.name in [\"vit_ti\", \"pit_ti\"]:  # for ViT: Drop CLS token\n",
    "    latents = [latent[:,1:] for latent in latents]\n",
    "latents = latents[:-1]  # drop logit (output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2799564d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "from einops import rearrange, reduce, repeat\n",
    "\n",
    "\n",
    "# aggregate feature map variances\n",
    "variances = []\n",
    "for latent in latents:  # `latents` is a list of hidden feature maps in latent spaces\n",
    "    latent = latent.cpu()\n",
    "    \n",
    "    if len(latent.shape) == 3:  # for ViT\n",
    "        b, n, c = latent.shape\n",
    "        h, w = int(math.sqrt(n)), int(math.sqrt(n))\n",
    "        latent = rearrange(latent, \"b (h w) c -> b c h w\", h=h, w=w)\n",
    "    elif len(latent.shape) == 4:  # for CNN\n",
    "        b, c, h, w = latent.shape\n",
    "    else:\n",
    "        raise Exception(\"shape: %s\" % str(latent.shape))\n",
    "                \n",
    "    variances.append(latent.var(dim=[-1, -2]).mean(dim=[0, 1]))\n",
    "    \n",
    "\n",
    "# Plot Fig 9: \"Feature map variance\"\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "if model.name == \"resnet_50\":  # for ResNet-50\n",
    "    pools = [4, 8, 14]\n",
    "    msas = []\n",
    "    marker = \"D\"\n",
    "    color = \"tab:blue\"\n",
    "elif model.name == \"vit_ti\":  # for ViT-Ti\n",
    "    pools = []\n",
    "    msas = [1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23,]\n",
    "    marker = \"o\"\n",
    "    color = \"tab:red\"\n",
    "else:\n",
    "    import warnings\n",
    "    warnings.warn(\"The configuration for %s are not implemented.\" % model.name, Warning)\n",
    "    pools, msas = [], []\n",
    "    marker = \"s\"\n",
    "    color = \"tab:green\"\n",
    "\n",
    "depths = range(len(variances))\n",
    "\n",
    "# normalize\n",
    "depth = len(depths) - 1\n",
    "depths = (np.array(depths)) / depth\n",
    "pools = (np.array(pools)) / depth\n",
    "msas = (np.array(msas)) / depth\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(6.5, 4), dpi=200)\n",
    "ax.plot(depths, variances, marker=marker, color=color, markersize=7)\n",
    "\n",
    "for pool in pools:\n",
    "    ax.axvspan(pool - 1.0 / depth, pool + 0.0 / depth, color=\"tab:blue\", alpha=0.15, lw=0)\n",
    "for msa in msas:\n",
    "    ax.axvspan(msa - 1.0 / depth, msa + 0.0 / depth, color=\"tab:gray\", alpha=0.15, lw=0)\n",
    "    \n",
    "ax.set_xlim(left=0, right=1.0)\n",
    "ax.set_ylim(bottom=0.0,)\n",
    "\n",
    "ax.set_xlabel(\"Normalized depth\")\n",
    "ax.set_ylabel(\"Feature map variance\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35a820cb",
   "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}