{"cells":[{"metadata":{},"cell_type":"markdown","source":"# About this notebook  \n- Visualize ViT Attention Map\n- ViT github is [here](https://github.com/tczhangzhi/VisionTransformer-Pytorch).\n(I modified a little for attention map. please see this [issue](https://github.com/tczhangzhi/VisionTransformer-Pytorch/issues/1#issuecomment-739138519).)\n\n\nI want to show that Attention Map for cassava.\n- I just show a few sample in 2019 train dataset.\n\nYou can check my pretrained ViT weight in [here](https://www.kaggle.com/piantic/cassava-vit-b-16).\n\n### If this kernel is useful, feel free to upvote:)"},{"metadata":{},"cell_type":"markdown","source":"# Vision Transformer (ViT) : Attention Map"},{"metadata":{},"cell_type":"markdown","source":"This is the Attention Map example.\n- Reference is [here](https://github.com/jeonsworld/ViT-pytorch/blob/main/visualize_attention_map.ipynb)."},{"metadata":{},"cell_type":"markdown","source":"<img src='https://user-images.githubusercontent.com/6073256/101206904-2a338f00-36b3-11eb-8920-f617abab1604.png'>"},{"metadata":{},"cell_type":"markdown","source":"Next, we will see the attention map for cassava leaf!"},{"metadata":{},"cell_type":"markdown","source":"# Import libraries"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import sys\n\npackage_path = '../input/visiontransformerpytorch121/VisionTransformer-Pytorch'\nsys.path.append(package_path)\n\nimport os\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport numpy as np\n\nimport torch\nimport torch.nn as nn\nimport torchvision.models as models\n\nimport cv2\nfrom PIL import Image\nfrom torchvision import transforms\n\nfrom vision_transformer_pytorch import VisionTransformer","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"transform = transforms.Compose([\n    transforms.Resize((384, 384)),\n    transforms.ToTensor(),\n    transforms.Normalize(\n        mean=[0.485, 0.456, 0.406],\n        std=[0.229, 0.224, 0.225],\n    ),\n])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Helper function"},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"# ====================================================\n# Helper functions\n# ====================================================\ndef load_state(model_path):\n    state_dict = torch.load(model_path)['model']\n    state_dict = {k[7:] if k.startswith('module.') else k: state_dict[k] for k in state_dict.keys()}\n    state_dict = {k[6:] if k.startswith('model.') else k: state_dict[k] for k in state_dict.keys()}\n\n    return state_dict","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"def get_attention_map(img, get_mask=False):\n    x = transform(img)\n    x.size()\n\n    logits, att_mat = model(x.unsqueeze(0))\n\n    att_mat = torch.stack(att_mat).squeeze(1)\n\n    # Average the attention weights across all heads.\n    att_mat = torch.mean(att_mat, dim=1)\n\n    # To account for residual connections, we add an identity matrix to the\n    # attention matrix and re-normalize the weights.\n    residual_att = torch.eye(att_mat.size(1))\n    aug_att_mat = att_mat + residual_att\n    aug_att_mat = aug_att_mat / aug_att_mat.sum(dim=-1).unsqueeze(-1)\n\n    # Recursively multiply the weight matrices\n    joint_attentions = torch.zeros(aug_att_mat.size())\n    joint_attentions[0] = aug_att_mat[0]\n\n    for n in range(1, aug_att_mat.size(0)):\n        joint_attentions[n] = torch.matmul(aug_att_mat[n], joint_attentions[n-1])\n\n    v = joint_attentions[-1]\n    grid_size = int(np.sqrt(aug_att_mat.size(-1)))\n    mask = v[0, 1:].reshape(grid_size, grid_size).detach().numpy()\n    if get_mask:\n        result = cv2.resize(mask / mask.max(), img.size)\n    else:        \n        mask = cv2.resize(mask / mask.max(), img.size)[..., np.newaxis]\n        result = (mask * img).astype(\"uint8\")\n    \n    return result\n\ndef plot_attention_map(original_img, att_map):\n    fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(16, 16))\n    ax1.set_title('Original')\n    ax2.set_title('Attention Map Last Layer')\n    _ = ax1.imshow(original_img)\n    _ = ax2.imshow(att_map)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Load ViT Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"model = VisionTransformer.from_name('ViT-B_16', num_classes=5)\nstate = load_state('../input/cassava-vit-b-16/ViT-B_16_fold0.pth')\nmodel.load_state_dict(state)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Visualize Attention Map"},{"metadata":{"trusted":true},"cell_type":"code","source":"label_map = pd.read_json('../input/cassava-leaf-disease-classification/label_num_to_disease_map.json', \n                         orient='index')\n\ndisplay(label_map)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## CBB - Class0"},{"metadata":{"trusted":true},"cell_type":"code","source":"img1 = Image.open(\"../input/cassava-vit-b-16/train-cbb-114.jpg\")\nimg2 = Image.open(\"../input/cassava-vit-b-16/train-cbb-44.jpg\")\n\nresult1 = get_attention_map(img1)\nresult2 = get_attention_map(img2)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img1, result1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img2, result2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Check mask for Attention Map"},{"metadata":{"trusted":true},"cell_type":"code","source":"result1 = get_attention_map(img1, True)\nresult2 = get_attention_map(img2, True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img1, result1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img2, result2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## CBSD - Class1"},{"metadata":{"trusted":true},"cell_type":"code","source":"img1 = Image.open(\"../input/cassava-vit-b-16/train-cbsd-154.jpg\")\nimg2 = Image.open(\"../input/cassava-vit-b-16/train-cbsd-821.jpg\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"result1 = get_attention_map(img1)\nresult2 = get_attention_map(img2)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img1, result1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img2, result2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Check mask for Attention Map"},{"metadata":{"trusted":true},"cell_type":"code","source":"result1 = get_attention_map(img1, True)\nresult2 = get_attention_map(img2, True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img1, result1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img2, result2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## CGM - Class2"},{"metadata":{"trusted":true},"cell_type":"code","source":"img1 = Image.open(\"../input/cassava-vit-b-16/train-cgm-498.jpg\")\nimg2 = Image.open(\"../input/cassava-vit-b-16/train-cgm-6.jpg\")\n\nresult1 = get_attention_map(img1)\nresult2 = get_attention_map(img2)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img1, result1)\nplot_attention_map(img2, result2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Check mask for Attention Map"},{"metadata":{"trusted":true},"cell_type":"code","source":"result1 = get_attention_map(img1, True)\nresult2 = get_attention_map(img2, True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img1, result1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img2, result2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## CMD - Class3"},{"metadata":{"trusted":true},"cell_type":"code","source":"img1 = Image.open(\"../input/cassava-vit-b-16/train-cbsd-154.jpg\")\nimg2 = Image.open(\"../input/cassava-vit-b-16/train-cbsd-821.jpg\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"result1 = get_attention_map(img1)\nresult2 = get_attention_map(img2)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img1, result1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img2, result2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Check mask for Attention Map"},{"metadata":{"trusted":true},"cell_type":"code","source":"result1 = get_attention_map(img1, True)\nresult2 = get_attention_map(img2, True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img1, result1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img2, result2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Healthy - Class4"},{"metadata":{"trusted":true},"cell_type":"code","source":"img1 = Image.open(\"../input/cassava-vit-b-16/train-healthy-105.jpg\")\nimg2 = Image.open(\"../input/cassava-vit-b-16/train-healthy-236.jpg\")\n\nresult1 = get_attention_map(img1)\nresult2 = get_attention_map(img2)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img1, result1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img2, result2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Check mask for Attention Map"},{"metadata":{"trusted":true},"cell_type":"code","source":"result1 = get_attention_map(img1, True)\nresult2 = get_attention_map(img2, True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img1, result1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_attention_map(img2, result2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Visualize Attention Maps"},{"metadata":{},"cell_type":"markdown","source":"For example, I will use cbsd images."},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"def get_attention_info(img):\n    x = transform(img)\n    x.size()\n\n    logits, att_mat = model(x.unsqueeze(0))\n\n    att_mat = torch.stack(att_mat).squeeze(1)\n\n    # Average the attention weights across all heads.\n    att_mat = torch.mean(att_mat, dim=1)\n\n    # To account for residual connections, we add an identity matrix to the\n    # attention matrix and re-normalize the weights.\n    residual_att = torch.eye(att_mat.size(1))\n    aug_att_mat = att_mat + residual_att\n    aug_att_mat = aug_att_mat / aug_att_mat.sum(dim=-1).unsqueeze(-1)\n\n    # Recursively multiply the weight matrices\n    joint_attentions = torch.zeros(aug_att_mat.size())\n    joint_attentions[0] = aug_att_mat[0]\n\n    for n in range(1, aug_att_mat.size(0)):\n        joint_attentions[n] = torch.matmul(aug_att_mat[n], joint_attentions[n-1])\n\n    grid_size = int(np.sqrt(aug_att_mat.size(-1)))\n    \n    return joint_attentions, grid_size\n\n# def plot_attention_map(original_img, att_map):\n#     fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(16, 16))\n#     ax1.set_title('Original')\n#     ax2.set_title('Attention Map Last Layer')\n#     _ = ax1.imshow(original_img)\n#     _ = ax2.imshow(att_map)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"img1 = Image.open(\"../input/cassava-vit-b-16/train-cbsd-154.jpg\")\nimg2 = Image.open(\"../input/cassava-vit-b-16/train-cbsd-821.jpg\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"joint_att1, grid_size1 = get_attention_info(img1)\njoint_att2, grid_size2 = get_attention_info(img2)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for i, v in enumerate(joint_att1):\n    v = joint_att1[-1]\n    mask = v[0, 1:].reshape(grid_size1, grid_size1).detach().numpy()\n    mask = cv2.resize(mask / mask.max(), img1.size)[..., np.newaxis]\n    result = (mask * img1).astype(\"uint8\")\n\n    fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(16, 16))\n    ax1.set_title('Original')\n    ax2.set_title('Attention Map_%d Layer' % (i+1))\n    _ = ax1.imshow(img1)\n    _ = ax2.imshow(result)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for i, v in enumerate(joint_att2):\n    v = joint_att2[-1]\n    mask = v[0, 1:].reshape(grid_size2, grid_size2).detach().numpy()\n    mask = cv2.resize(mask / mask.max(), img1.size)[..., np.newaxis]\n    result = (mask * img2).astype(\"uint8\")\n\n    fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(16, 16))\n    ax1.set_title('Original')\n    ax2.set_title('Attention Map_%d Layer' % (i+1))\n    _ = ax1.imshow(img2)\n    _ = ax2.imshow(result)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## If this kernel is useful, <font color='orange'>please upvote</font>!"}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}