token_visualization.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import torchvision.transforms as transforms\n",
    "import torch.nn.functional as F\n",
    "import torch\n",
    "import adaptive_tokenizers\n",
    "from utils import misc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For all token visualizations in the paper, we used alit_small_vqgan_quantized_latents.pth model (attention maps from the 6th layer of the decoder).\n",
    "# Interesting token-object binding visualizations are on images with more than one object, eg: coco images (even though the model is trained only on Imagenet100)\n",
    "# The last iteration attention maps are most crisp and sharp, and have best token-object binding. This is because the presence of more tokens (or memory) in the last iteration, allows some tokens to specialize and attend to objects.\n",
    "# Try visualizing the few coco images we have attached in assets/ directory.\n",
    "\n",
    "args = {\n",
    "    'image_path': 'assets/custom_images/coco/val2017_000000000632.jpg',\n",
    "    'device': 'cuda:1',\n",
    "    'input_size': 256,\n",
    "    'model': 'alit_small',\n",
    "    'base_tokenizer': 'vqgan',\n",
    "    'ckpt': 'adaptive_tokenizers/pretrained_models/imagenet100/alit_small_vqgan_quantized_latents.pth',\n",
    "    'quantize_latent': True, # 1D quantized models have better object-token binding.\n",
    "    'num_layers': 8, # 8 for small, 12 for base, 16 for large\n",
    "    'layer_to_visualize': 6,\n",
    "    'iter_to_visualize': 7 # 0<=iter_to_visualize<=7\n",
    "}\n",
    "args = misc.Args(**args)\n",
    "\n",
    "image = Image.open(args.image_path).convert(\"RGB\")\n",
    "img = np.asarray(image)\n",
    "transform_val = transforms.Compose([\n",
    "    transforms.Resize(args.input_size, interpolation=transforms.InterpolationMode.BICUBIC, antialias=True),\n",
    "    transforms.CenterCrop(args.input_size),\n",
    "    transforms.ToTensor()\n",
    "])\n",
    "image_tensor = transform_val(image).to(args.device)[None]\n",
    "\n",
    "plt.imshow(image_tensor[0].permute([1,2,0]).cpu().numpy())\n",
    "plt.show()\n",
    "\n",
    "base_tokenizer_args = {\n",
    "    \"id\": args.base_tokenizer,\n",
    "    \"is_requires_grad\": False\n",
    "}\n",
    "adaptive_tokenizer = adaptive_tokenizers.__dict__[args.model](\n",
    "    base_tokenizer_args=base_tokenizer_args, quantize_latent=args.quantize_latent, \n",
    "    train_stage=\"full_finetuning\", visualize_decoder_attn_weights=True)\n",
    "\n",
    "adaptive_tokenizer.to(args.device)\n",
    "checkpoint = torch.load(args.ckpt, map_location='cpu')\n",
    "adaptive_tokenizer.load_state_dict(checkpoint['ema'], strict=True)\n",
    "adaptive_tokenizer.eval()\n",
    "\n",
    "with torch.no_grad():\n",
    "    # Automatic sample minimum length representation for the image.\n",
    "    # Currently, we support only \"Reconstruction Loss < Threshold\" as automatic Token Selection Criteria (TSC).\n",
    "    reconstruction_loss_threshold = 0.05\n",
    "    _, _, all_logs = adaptive_tokenizer.encode(image_tensor, return_min_length_embedding=False, token_selection_criteria=\"reconstruction_loss\", threshold=reconstruction_loss_threshold, return_embedding_type=\"latent_tokens\")\n",
    "    \n",
    "    num_tokens = 32\n",
    "    start_tokens = 0\n",
    "    \n",
    "    for iter in range(8):\n",
    "\n",
    "        fig, axes = plt.subplots(4, 8, figsize=(20, 10))\n",
    "        plt.subplots_adjust(wspace=0, hspace=0)\n",
    "        for idx in range(start_tokens, start_tokens + num_tokens):\n",
    "            attn_weight = np.concatenate(all_logs[iter][\"decoded_attn_weights_{}\".format(iter)])\n",
    "            attn_weight_vis = attn_weight[args.layer_to_visualize, 1:257, 256 + idx].reshape(16, 16)\n",
    "            attn_weight_vis = torch.tensor(attn_weight_vis).unsqueeze(0).unsqueeze(0)  # Shape (1, 1, 16, 16)\n",
    "            attn_weight_vis = F.interpolate(attn_weight_vis, size=(256, 256), mode='bilinear')\n",
    "            attn_weight_vis = attn_weight_vis.squeeze()  # Shape (256, 256)\n",
    "            unnormalized_high_attention_mask = attn_weight_vis\n",
    "\n",
    "            row = (idx - start_tokens) // 8  # 4 rows for 32 tokens\n",
    "            col = (idx - start_tokens) % 8   # 8 columns for 32 tokens\n",
    "\n",
    "            min_val, max_val = unnormalized_high_attention_mask.min(), unnormalized_high_attention_mask.max() \n",
    "            # max_val = 0.03 # for vae based model, setting optimal max_val leads to better visualization\n",
    "            high_attention_mask = (unnormalized_high_attention_mask - min_val) / (max_val - min_val)\n",
    "            high_attention_mask[high_attention_mask>1]=1\n",
    "            im = axes[row, col].imshow(high_attention_mask, cmap='jet', aspect='auto')\n",
    "            axes[row, col].imshow(img, alpha=0.4, aspect='auto')\n",
    "            axes[row, col].axis('off')\n",
    "        \n",
    "        plt.show(fig)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "adaptive_representations",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}