notebooks/mean-attention-distance-1k.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sR8zW9qDql2q"
   },
   "source": [
    "## Setup\n",
    "\n",
    "In this notebook we compute mean attention distances over 1000 data points from ImageNet-1k validation set. Our main source of reference is: [Do Vision Transformers See Like Convolutional Neural Networks?](https://arxiv.org/abs/2108.08810)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DZmljdmmrMtV",
    "outputId": "6e52f08f-f85b-4575-efc3-2f99b3d95d24"
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade gdown -qq\n",
    "!gdown --id 1oeukDq54YMV5xWFMq0j1VN25HQQAd9ei\n",
    "!unzip -q 1000_val_images_sampled.zip\n",
    "!wget -q https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt -O ilsvrc2012_wordnet_lemmas.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lpMhcpWWPdt3"
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2rA5vozUqf2_"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_hub as hub\n",
    "from tensorflow import keras\n",
    "\n",
    "from PIL import Image\n",
    "from io import BytesIO\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import requests\n",
    "import zipfile\n",
    "import gdown\n",
    "import cv2\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NMDOuhKzPhX1"
   },
   "source": [
    "## Chose the ViT model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_JEzN45BrEJm",
    "outputId": "9d46993e-5621-417f-b0c1-9cd86533879d"
   },
   "outputs": [],
   "source": [
    "model_name = \"vit_base_i21k_patch16_224\"\n",
    "\n",
    "model_handle_map = {\n",
    "    \"vit_base_i1k_patch16_224\": \"1mbtnliT3jRb3yJUHhbItWw8unfYZw8KJ\",\n",
    "    \"vit_base_i21k_patch16_224\": \"1ApOdYe4NXxhPhJABefgZ3KVvqsQzhCL7\",\n",
    "    \"dino_base_patch16_224\": \"16_1oDm0PeCGJ_KGBG5UKVN7TsAtiRNrN\",\n",
    "    \"deit_base_patch16_224\": \"https://tfhub.dev/sayakpaul/deit_base_patch16_224/1\",\n",
    "    \"deit_base_distilled_patch16_224\": \"https://tfhub.dev/sayakpaul/deit_base_distilled_patch16_224/1\",\n",
    "}\n",
    "\n",
    "# Derive the patch size, image resolution, and class tokens from the model name.\n",
    "splits = model_name.split(\"_\")\n",
    "model_type = splits[0]\n",
    "input_resolution = int(splits[-1])\n",
    "patch_size = int(splits[-2].replace(\"patch\", \"\"))\n",
    "num_cls_tokens = 2 if \"distilled\" in model_name else 1\n",
    "\n",
    "# Get the model handle.\n",
    "model_handle = model_handle_map[model_name]\n",
    "\n",
    "print(f\"Patch Size: {patch_size}.\")\n",
    "print(f\"Input resolution: {input_resolution} x {input_resolution} x 3.\")\n",
    "print(f\"TF-Hub handle OR Drive ID: {model_handle}.\")\n",
    "print(f\"Number of class tokens: {num_cls_tokens}.\")\n",
    "print(f\"Model type: {model_type}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AUrsuNs5RSYb"
   },
   "source": [
    "## Image preprocessing utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UDvOqRj3rFKj"
   },
   "outputs": [],
   "source": [
    "crop_layer = keras.layers.CenterCrop(input_resolution, input_resolution)\n",
    "norm_layer = keras.layers.Normalization(\n",
    "    mean=[0.485 * 255, 0.456 * 255, 0.406 * 255],\n",
    "    variance=[(0.229 * 255) ** 2, (0.224 * 255) ** 2, (0.225 * 255) ** 2],\n",
    ")\n",
    "rescale_layer = keras.layers.Rescaling(scale=1.0 / 127.5, offset=-1)\n",
    "\n",
    "\n",
    "def preprocess_image(image, size=input_resolution):\n",
    "    # turn the image into a numpy array and add batch dim\n",
    "    image = tf.expand_dims(image, 0)\n",
    "\n",
    "    # if model type is vit rescale the image to [-1, 1]\n",
    "    if model_type == \"vit\":\n",
    "        image = rescale_layer(image)\n",
    "\n",
    "    # resize the image using bicubic interpolation\n",
    "    resize_size = int((256 / 224) * size)\n",
    "    image = tf.image.resize(image, (resize_size, resize_size), method=\"bicubic\")\n",
    "\n",
    "    # crop the image\n",
    "    image = crop_layer(image)\n",
    "\n",
    "    # if model type is deit normalize the image\n",
    "    if model_type != \"vit\":\n",
    "        image = norm_layer(image)\n",
    "\n",
    "    # return the image\n",
    "    return image.numpy()\n",
    "\n",
    "\n",
    "def load_image_from_path(path):\n",
    "    image = Image.open(path)\n",
    "    image = np.array(image)\n",
    "    if len(image.shape) < 3:\n",
    "        image = np.tile(image[..., None], (1, 1, 3))\n",
    "    preprocessed_image = preprocess_image(image)\n",
    "    return image, preprocessed_image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SEym7PSQRbcM"
   },
   "source": [
    "## Prepare `tf.data.Dataset` object and load a model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Lrvs9JVHrQ65"
   },
   "outputs": [],
   "source": [
    "def get_tfhub_model(model_url: str) -> tf.keras.Model:\n",
    "    inputs = tf.keras.Input((input_resolution, input_resolution, 3))\n",
    "    hub_module = hub.KerasLayer(model_url)\n",
    "    outputs, attention_weights = hub_module(inputs)\n",
    "    return tf.keras.Model(inputs, outputs=[outputs, attention_weights])\n",
    "\n",
    "\n",
    "def get_gdrive_model(model_id: str) -> tf.keras.Model:\n",
    "    model_path = gdown.download(id=model_id, quiet=False)\n",
    "    with zipfile.ZipFile(model_path, \"r\") as zip_ref:\n",
    "        zip_ref.extractall()\n",
    "    model_name = model_path.split(\".\")[0]\n",
    "    inputs = tf.keras.Input((input_resolution, input_resolution, 3))\n",
    "    model = tf.keras.models.load_model(model_name, compile=False)\n",
    "    outputs, attention_weights = model(inputs)\n",
    "    return tf.keras.Model(inputs, outputs=[outputs, attention_weights])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "H47NQFiirS5Z",
    "outputId": "ac752c49-4f8c-4273-ea03-82683e09aef3"
   },
   "outputs": [],
   "source": [
    "if len(model_handle.split(\"/\")) > 1:\n",
    "    loaded_model = get_tfhub_model(model_handle)\n",
    "else:\n",
    "    loaded_model = get_gdrive_model(model_handle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "1CM2X9Pj2Urd",
    "outputId": "0f640eaf-aeba-4704-8e47-c66bfcaecddc"
   },
   "outputs": [],
   "source": [
    "list_images = os.listdir(\"1000_val_images_sampled\")\n",
    "list_images = [os.path.join(\"1000_val_images_sampled\", x) for x in list_images]\n",
    "\n",
    "images = list()\n",
    "preprocessed_images = list()\n",
    "for image_path in tqdm(list_images):\n",
    "    image, preprocessed_image = load_image_from_path(image_path)\n",
    "    images.append(image)\n",
    "    preprocessed_images.append(preprocessed_image[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NwZkF36m2WmA"
   },
   "outputs": [],
   "source": [
    "image_ds = (\n",
    "    tf.data.Dataset.from_tensor_slices(preprocessed_images)\n",
    "    .batch(200)\n",
    "    .prefetch(tf.data.AUTOTUNE)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HVWwz60ZRrIS"
   },
   "source": [
    "## Mean Attention Distance\n",
    "\n",
    "Reference: https://gist.github.com/simonster/155894d48aef2bd36bd2dd8267e62391\n",
    "\n",
    "### Compute Mean Distance\n",
    "![Compute Mean Distance](https://i.imgur.com/nyuS9H9.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gW4GkITDrdw1"
   },
   "outputs": [],
   "source": [
    "def compute_distance_matrix(patch_size, num_patches, length):\n",
    "    \"\"\"Helper function to compute distance matrix.\"\"\"\n",
    "    distance_matrix = np.zeros((num_patches, num_patches))\n",
    "    for i in range(num_patches):\n",
    "        for j in range(num_patches):\n",
    "            if i == j:  # zero distance\n",
    "                continue\n",
    "\n",
    "            xi, yi = (int(i / length)), (i % length)\n",
    "            xj, yj = (int(j / length)), (j % length)\n",
    "            distance_matrix[i, j] = patch_size * np.linalg.norm([xi - xj, yi - yj])\n",
    "\n",
    "    return distance_matrix\n",
    "\n",
    "\n",
    "def compute_mean_attention_dist(patch_size, attention_weights):\n",
    "    # The attention_weights shape = (batch, num_heads, num_patches, num_patches)\n",
    "    attention_weights = attention_weights[\n",
    "        ..., num_cls_tokens:, num_cls_tokens:\n",
    "    ]  # Removing the CLS token\n",
    "    num_patches = attention_weights.shape[-1]\n",
    "    length = int(np.sqrt(num_patches))\n",
    "    assert length ** 2 == num_patches, \"Num patches is not perfect square\"\n",
    "\n",
    "    distance_matrix = compute_distance_matrix(patch_size, num_patches, length)\n",
    "    h, w = distance_matrix.shape\n",
    "\n",
    "    distance_matrix = distance_matrix.reshape((1, 1, h, w))\n",
    "    # The attention_weights along the last axis adds to 1\n",
    "    # this is due to the fact that they are softmax of the raw logits\n",
    "    # summation of the (attention_weights * distance_matrix)\n",
    "    # should result in an average distance per token\n",
    "    mean_distances = attention_weights * distance_matrix\n",
    "    mean_distances = np.sum(\n",
    "        mean_distances, axis=-1\n",
    "    )  # sum along last axis to get average distance per token\n",
    "    mean_distances = np.mean(\n",
    "        mean_distances, axis=-1\n",
    "    )  # now average across all the tokens\n",
    "\n",
    "    return mean_distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_J82XUcqIokm",
    "outputId": "87a17f40-9677-450c-a971-226a17b31c36"
   },
   "outputs": [],
   "source": [
    "mean_distances = dict()\n",
    "for idx, image in tqdm(enumerate(image_ds)):\n",
    "    _, attention_score_dict = loaded_model.predict(image)\n",
    "    # Calculate the mean distances for every transformer block.\n",
    "    for name, attention_weight in attention_score_dict.items():\n",
    "        mean_distance = compute_mean_attention_dist(\n",
    "            patch_size=patch_size,\n",
    "            attention_weights=attention_weight,\n",
    "        )\n",
    "        if idx == 0:\n",
    "            mean_distances[f\"{name}_mean_dist\"] = mean_distance\n",
    "        else:\n",
    "            mean_distances[f\"{name}_mean_dist\"] = tf.concat(\n",
    "                [mean_distance, mean_distances[f\"{name}_mean_dist\"]], axis=0\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "m05ff8YEkZzy",
    "outputId": "ddf0015d-ab8e-417d-c476-a92fa5b6382c"
   },
   "outputs": [],
   "source": [
    "# For a single transforer block, we have attention distances\n",
    "# for 1000 images for each attention head. We have 12 such\n",
    "# attention head per transformer block.\n",
    "mean_distances[\"transformer_block_0_att_mean_dist\"].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DVhPXsnYSyQF",
    "outputId": "6e8abad9-bbad-451e-f261-5ef8fd1041b0"
   },
   "outputs": [],
   "source": [
    "# Get the number of heads from the mean distance output\n",
    "num_heads = tf.shape(mean_distances[\"transformer_block_0_att_mean_dist\"])[-1].numpy()\n",
    "\n",
    "# Print the shapes\n",
    "print(f\"Num Heads: {num_heads}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6PKZvtDKK0AC"
   },
   "outputs": [],
   "source": [
    "for key, value in mean_distances.items():\n",
    "    mean_distances[key] = tf.reduce_mean(value, axis=0, keepdims=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "f7wjXL7yk_F5",
    "outputId": "c03016c9-702c-47df-a081-f8c051ca539e"
   },
   "outputs": [],
   "source": [
    "# For a single transformer block we have calculate mean attention distances\n",
    "# for 12 attention heads.\n",
    "mean_distances[\"transformer_block_0_att_mean_dist\"].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iqNJuPs7SKyT"
   },
   "source": [
    "## Visualize the mean distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 621
    },
    "id": "BpfokYqwzYzI",
    "outputId": "eb0c46dd-e251-4f8d-b43b-f0e4aa24f2e3"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 10))\n",
    "\n",
    "first = list()\n",
    "last = list()\n",
    "\n",
    "for idx in range(len(mean_distances)):\n",
    "    mean_distance = mean_distances[f\"transformer_block_{idx}_att_mean_dist\"]\n",
    "\n",
    "    x = [idx] * num_heads\n",
    "    y = mean_distance[0, :]\n",
    "\n",
    "    plt.scatter(x=x, y=y, label=f\"block_{idx}\")\n",
    "    for i, txt in enumerate(range(num_heads)):\n",
    "        plt.annotate(txt, (x[i] + 0.1, y[i] + 0.1))\n",
    "\n",
    "plt.xlabel(\"Transformer Blocks\")\n",
    "plt.ylabel(\"Mean Attention Distance\")\n",
    "plt.legend(loc=\"lower right\")\n",
    "plt.title(model_name)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 621
    },
    "id": "3B86CgwQso74",
    "outputId": "366a98e0-24c1-4ef8-c8b0-a853260e693e"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 10))\n",
    "\n",
    "first = list() \n",
    "last = list()\n",
    "\n",
    "for idx in range(len(mean_distances)):\n",
    "    mean_distance = mean_distances[f\"transformer_block_{idx}_att_mean_dist\"]\n",
    "\n",
    "    x = [idx] * num_heads\n",
    "    y = mean_distance[0, :]\n",
    "\n",
    "    plt.scatter(x=x, y=y, label=f\"block_{idx}\")\n",
    "\n",
    "    first.append(y[0])\n",
    "    last.append(y[-1])\n",
    "\n",
    "plt.plot(first, \"r-\", label=\"first head\")\n",
    "plt.plot(last, \"g-\", label=\"last head\")\n",
    "\n",
    "plt.xlabel(\"Transformer Blocks\")\n",
    "plt.ylabel(\"Mean Attention Distance\")\n",
    "plt.legend(loc=\"lower right\")\n",
    "plt.title(model_name)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 621
    },
    "id": "PLKgA360t0QZ",
    "outputId": "1e3686d2-2a7f-41e2-caee-0cd637ec4ac1"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 10))\n",
    "\n",
    "first = list()\n",
    "last = list()\n",
    "\n",
    "blocks = [4, 5, 10, 11]\n",
    "chars = [\"-\", \"--\", \"-.\", \":\"]\n",
    "colors = [\"b\", \"g\", \"r\", \"c\"]\n",
    "\n",
    "for idx, block_num in enumerate(blocks):\n",
    "    mean_distance = mean_distances[f\"transformer_block_{block_num}_att_mean_dist\"]\n",
    "\n",
    "    x = list(range(num_heads))\n",
    "    y = mean_distance[0, :]\n",
    "\n",
    "    plt.plot(x, y, f\"{colors[idx]}{chars[idx]}\", label=f\"block_{block_num}\")\n",
    "\n",
    "plt.xlabel(\"Transformer Heads\")\n",
    "plt.ylabel(\"Mean Attention Distance\")\n",
    "plt.legend(loc=\"lower right\")\n",
    "plt.title(model_name)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "machine_shape": "hm",
   "name": "probing1000.ipynb",
   "provenance": []
  },
  "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}