notebooks/single-instance-probing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sR8zW9qDql2q"
   },
   "source": [
    "## Introduction\n",
    "\n",
    "In this notebook we probe into the ViT attention weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DZmljdmmrMtV",
    "outputId": "633c20db-bc62-4931-c4e5-8285cfe22b68"
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade gdown -qq\n",
    "!wget https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt -O ilsvrc2012_wordnet_lemmas.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lpMhcpWWPdt3"
   },
   "source": [
    "## Imports and Setups"
   ]
  },
  {
   "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",
    "import numpy as np\n",
    "import requests\n",
    "import zipfile\n",
    "import gdown\n",
    "import cv2"
   ]
  },
  {
   "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": "f1fff414-8f74-4617-b334-866ab0d5db99"
   },
   "outputs": [],
   "source": [
    "model_name = \"dino_base_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": [
    "## Preprocessing the test image"
   ]
  },
  {
   "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./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 = np.array(image)\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(\n",
    "        image,\n",
    "        (resize_size, resize_size),\n",
    "        method=\"bicubic\"\n",
    "    )\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_url(url):\n",
    "    # Credit: Willi Gierke\n",
    "    response = requests.get(url)\n",
    "    image = Image.open(BytesIO(response.content))\n",
    "    preprocessed_image = preprocess_image(image)\n",
    "    return image, preprocessed_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 248
    },
    "id": "A16LHAIHrHY0",
    "outputId": "98016d7e-391d-4170-c520-d97651059e00"
   },
   "outputs": [],
   "source": [
    "with open(\"ilsvrc2012_wordnet_lemmas.txt\", \"r\") as f:\n",
    "    lines = f.readlines()\n",
    "imagenet_int_to_str = [line.rstrip() for line in lines]\n",
    "\n",
    "img_url = \"https://p0.pikrepo.com/preview/853/907/close-up-photo-of-gray-elephant.jpg\"\n",
    "image, preprocessed_image = load_image_from_url(img_url)\n",
    "\n",
    "plt.imshow(image)\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SEym7PSQRbcM"
   },
   "source": [
    "## Run inference"
   ]
  },
  {
   "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": "13507add-2974-4908-ebaf-1c4fac05359d"
   },
   "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": "MvfpQeJb5Nnb",
    "outputId": "f643026b-60e9-43c2-c6b4-574761c67868"
   },
   "outputs": [],
   "source": [
    "predictions, attention_score_dict = loaded_model.predict(preprocessed_image)\n",
    "predicted_label = imagenet_int_to_str[int(np.argmax(predictions))]\n",
    "print(predicted_label)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Bj00PxF4RiWX"
   },
   "source": [
    "## Probing into the attention weights\n",
    "\n",
    "Reference: https://colab.research.google.com/github/sayakpaul/deit-tf/blob/main/notebooks/classification.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Fnyv_NE-rY1r",
    "outputId": "95ba6907-2a47-4a17-f45f-346398cb4841"
   },
   "outputs": [],
   "source": [
    "# Stack the individual attention matrices from individual transformer blocks.\n",
    "attn_mat = tf.stack([attention_score_dict[k] for k in attention_score_dict.keys()])\n",
    "attn_mat = tf.squeeze(attn_mat, axis=1)\n",
    "print(attn_mat.shape)\n",
    "\n",
    "# Average the attention weights across all heads.\n",
    "attn_mat = tf.reduce_mean(attn_mat, axis=1)\n",
    "print(attn_mat.shape)\n",
    "\n",
    "# To account for residual connections, we add an identity matrix to the\n",
    "# attention matrix and re-normalize the weights.\n",
    "residual_attn = tf.eye(attn_mat.shape[1])\n",
    "aug_attn_mat = attn_mat + residual_attn\n",
    "aug_attn_mat = aug_attn_mat / tf.reduce_sum(aug_attn_mat, axis=-1)[..., None]\n",
    "aug_attn_mat = aug_attn_mat.numpy()\n",
    "print(aug_attn_mat.shape)\n",
    "\n",
    "# Recursively multiply the weight matrices\n",
    "joint_attentions = np.zeros(aug_attn_mat.shape)\n",
    "joint_attentions[0] = aug_attn_mat[0]\n",
    "\n",
    "for n in range(1, aug_attn_mat.shape[0]):\n",
    "    joint_attentions[n] = np.matmul(aug_attn_mat[n], joint_attentions[n-1])\n",
    "    \n",
    "# Attention from the output token to the input space.\n",
    "v = joint_attentions[-1]\n",
    "grid_size = int(np.sqrt(aug_attn_mat.shape[-1]))\n",
    "mask = v[0, num_cls_tokens:].reshape(grid_size, grid_size)\n",
    "mask = cv2.resize(mask / mask.max(), image.size)[..., np.newaxis]\n",
    "result = (mask * image).astype(\"uint8\")\n",
    "print(result.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NGne7rQGRnyX"
   },
   "source": [
    "## Visualize the attention maps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 316
    },
    "id": "ZoXwTajkrayd",
    "outputId": "afbeda84-4905-421e-f668-e048bc0d74b3"
   },
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(10, 8))\n",
    "fig.suptitle(f\"Predicted label: {predicted_label}.\", fontsize=20)\n",
    "\n",
    "_ = ax1.imshow(image)\n",
    "_ = ax2.imshow(result)\n",
    "ax1.set_title(\"Input Image\", fontsize=16)\n",
    "ax2.set_title(\"Attention Map\", fontsize=16)\n",
    "ax1.axis(\"off\")\n",
    "ax2.axis(\"off\")\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.subplots_adjust(top=1.35)\n",
    "fig.show()"
   ]
  },
  {
   "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",
    "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[..., num_cls_tokens:, num_cls_tokens:] # 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(mean_distances, axis=-1) # sum along last axis to get average distance per token\n",
    "    mean_distances = np.mean(mean_distances, axis=-1) # now average across all the tokens\n",
    "\n",
    "    return mean_distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "S86mLJwusM1d"
   },
   "outputs": [],
   "source": [
    "# Build the mean distances for every transformer block\n",
    "mean_distances = {\n",
    "    f\"{name}_mean_dist\": compute_mean_attention_dist(patch_size=patch_size, attention_weights=attention_weight)\n",
    "    for name, attention_weight in attention_score_dict.items()\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DVhPXsnYSyQF",
    "outputId": "551f5b5e-6c15-4644-98f3-dc18967f8ba6"
   },
   "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": "markdown",
   "metadata": {
    "id": "iqNJuPs7SKyT"
   },
   "source": [
    "## Visualize the mean distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 607
    },
    "id": "BpfokYqwzYzI",
    "outputId": "e52739e1-dc00-49ec-d600-c819109986a4"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 10))\n",
    "\n",
    "for idx in range(len(mean_distances)):\n",
    "    mean_distance = mean_distances[f\"transformer_block_{idx}_att_mean_dist\"]\n",
    "    x = [idx] * num_heads\n",
    "    y = mean_distance[0, :]\n",
    "    plt.scatter(x=x, y=y, label=f\"attention_{idx}\")\n",
    "\n",
    "plt.legend(loc=\"lower right\")\n",
    "plt.title(model_name)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "machine_shape": "hm",
   "name": "probing-thought-exp.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
}