Visualize attention map for vision transformer

[kiashann]

Hi, I want to extract attention map from pretrained vision transformer for specific image.
How I can do that?

[hankyul2]

Hi @kiashann

This is toy examples to visualize whole attention map and attention map only for class token. (see here for more information)

import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from timm.models import create_model
import torch.nn.functional as F
from torchvision.transforms import Compose, Resize, CenterCrop, Normalize, ToTensor

def to_tensor(img):
    transform_fn = Compose([Resize(249, 3), CenterCrop(224), ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    return transform_fn(img)

def show_img(img):
    img = np.asarray(img)
    plt.figure(figsize=(10, 10))
    plt.imshow(img)
    plt.axis('off')
    plt.show()

def show_img2(img1, img2, alpha=0.8):
    img1 = np.asarray(img1)
    img2 = np.asarray(img2)
    plt.figure(figsize=(10, 10))
    plt.imshow(img1)
    plt.imshow(img2, alpha=alpha)
    plt.axis('off')
    plt.show()

def my_forward_wrapper(attn_obj):
    def my_forward(x):
        B, N, C = x.shape
        qkv = attn_obj.qkv(x).reshape(B, N, 3, attn_obj.num_heads, C // attn_obj.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)   # make torchscript happy (cannot use tensor as tuple)

        attn = (q @ k.transpose(-2, -1)) * attn_obj.scale
        attn = attn.softmax(dim=-1)
        attn = attn_obj.attn_drop(attn)
        attn_obj.attn_map = attn
        attn_obj.cls_attn_map = attn[:, :, 0, 2:]

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = attn_obj.proj(x)
        x = attn_obj.proj_drop(x)
        return x
    return my_forward

img = Image.open('n02102480_Sussex_spaniel.JPEG')
x = to_tensor(img)

model = create_model('deit_small_distilled_patch16_224', pretrained=True)
model.blocks[-1].attn.forward = my_forward_wrapper(model.blocks[-1].attn)

y = model(x.unsqueeze(0))
attn_map = model.blocks[-1].attn.attn_map.mean(dim=1).squeeze(0).detach()
cls_weight = model.blocks[-1].attn.cls_attn_map.mean(dim=1).view(14, 14).detach()

img_resized = x.permute(1, 2, 0) * 0.5 + 0.5
cls_resized = F.interpolate(cls_weight.view(1, 1, 14, 14), (224, 224), mode='bilinear').view(224, 224, 1)

show_img(img)
show_img(attn_map)
show_img(cls_weight)
show_img(img_resized)
show_img2(img_resized, cls_resized, alpha=0.8)

original image

attention map for last layer (198 x 198 (=196(img) + 1(cls) + 1(distill)))

class attention map for last layer (14 x 14)

class attention map over image

[mae338]

@kiashann
Thank you for your valuable code. the whole code is working fine but I just need to understand how these lines work : model.blocks[-1].attn.forward=my_forward_wrapper(model.blocks[-1].attn) and attn_map = model.blocks[-1].attn.attn_map.mean(dim=1).squeeze(0).detach() so I ran these lines of code in the console : attn_obj=model.blocks[-1].attn & qkv = attn_obj.qkv(x) but got this error (RuntimeError: mat1 and mat2 shapes cannot be multiplied (672x224 and 384x1152)). I'd like to know why. I need to make sure whether x is the transformed image or some other variable – When I debugged the code I found out that B,N,C from x.shape are 1 &384 &198, which are different from the dimensions of the transformed image

[hankyul2]

Hi @mae338

I am happy to help you.

• In y=model(x.unsqueeze(0)), x.unsqueeze(0) is a transformed image that has shapes as (1, 3, 224, 224).
• In model.blocks[-1].attn.forward=my_forward_wrapper(model.blocks[-1].attn), we replace the original forward to our my_forward_wrapper to save attention map as an instance variable of model.blocks[-1].attn.
• In model.blocks[-1].attn.attn_map.mean(dim=1).squeeze(0).detach(), we average attention map in head-dimension .mean(dim=1) so that we can display whole attention.
• To understand better, we recommend you print out every shape of the tensor, which helps you to see the overall workflow of ViT.

Thank you.

Hankyul

[mae338]

Thank you so much, Hankyul. I appreciate your effort to help me. I'd be pleased if you could help me understand this as well:
What about def my_forward(x):
B, N, C = x.shape
qkv = attn_obj.qkv(x)?
more especifically, the x parameter? what does it refer to?
and why did you create the function my_forward?
I also found out that the function my_forward must return x or else it doesn't work and the whole program won't work. I wonder why?

[hankyul2]

Hi @mae338

I hope this could help you.

• We replace the original Attention forward() method to my_forward() method that only inserts two extra codes for saving attention map (attn_obj.attn_map = attn, attn_obj.cls_attn_map = attn[:, :, 0, 2:]). Thus, everything (function, signature, and return value) should be the same as the original function except for additional code.
• x is image tokens (Batch x N x D) input to attention. Image tokens are tokenized by a linear layer, and their shape depends on the model's size. In our case, DeiT-small/16 splits the whole image into 196 (N=198, with 2 extra class, distill tokens) patches and has 384 (D=384) channel dimensions.
• Since the attention block passes the output to the next block such as the MLP block, my_forward() should also pass the return value. If you skip passing the return value to the next block, the next blocks get a None value, which is an unexpected situation for them, thereby generating an error.

Thank you.

Hankyul

[mae338]

Hi @hankyul2,

I’d like to view the outputs of a pretrained vit model (vit-base-batch-16 224) , especially the input to the mlp head. I tried the same function but it didn’t work. Any suggestions?
Thank you
mae

[hankyul2]

Hi @mae338

I hope this can help you.

You can extract the ultimate features of a pre-trained ViT by y = model.forward_head(x, pre_logits=True) and visualize them for your purpose. Since it was a long time ago when I upload initial code in first comment, I copied and modified them as:

# dependency
!wget https://user-images.githubusercontent.com/31476895/167238573-b0cc3a6d-d3ee-462b-8630-a8f253e69bb2.png
!pip install -Uq fastai timm==0.6.13 huggingface_hub
##########################################

# code
import numpy as np
from PIL import Image
from timm.models import create_model
from torch import nn
from torchvision.transforms import Compose, Resize, CenterCrop, Normalize, ToTensor

def to_tensor(img):
    transform_fn = Compose([Resize(249), CenterCrop(224), ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    return transform_fn(img)

img = Image.open('167238573-b0cc3a6d-d3ee-462b-8630-a8f253e69bb2.png').convert('RGB')
x = to_tensor(img)

model = create_model('deit_small_distilled_patch16_224', pretrained=True)
x = model.forward_features(x.unsqueeze(0))
y = model.forward_head(x, pre_logits=True)
print(y.shape)
##########################################

# output
torch.Size([1, 384])
##########################################

Thank you.

Hankyul

[tomos7231]

@hankyul2

I would like to apply this code to the 'vit_small_patch16_384' model from timm. How should I modify the code for this purpose?
(I understand that '224' in the given code refers to the image size, but how is '14' determined?)
I apologize if this is due to my lack of knowledge.

attn_map = model.blocks[-1].attn.attn_map.mean(dim=1).squeeze(0).detach()
cls_weight = model.blocks[-1].attn.cls_attn_map.mean(dim=1).view(14, 14).detach()

img_resized = x.permute(1, 2, 0) * 0.5 + 0.5
cls_resized = F.interpolate(cls_weight.view(1, 1, 14, 14), (224, 224), mode='bilinear').view(224, 224, 1)

[hankyul2]

Hi @tomos7231

1. 384 means input image resolution to ViT model. If you want to extract an attention map using the code above, you should change the input resolution (224 -> 384) and the dimension of patches (14x14 -> 24x24) like below code.

attn_map = model.blocks[-1].attn.attn_map.mean(dim=1).squeeze(0).detach()
cls_weight = model.blocks[-1].attn.cls_attn_map.mean(dim=1).view(24, 24).detach() # (14->24)

img_resized = x.permute(1, 2, 0) * 0.5 + 0.5
cls_resized = F.interpolate(cls_weight.view(1, 1, 24, 24), (384, 384), mode='bilinear').view(384, 384, 1) # (14->24), (224->384)

2. 14 means the number of patches in each spatial dimension, e.g., H, W of the feature map. This value (14) is determined by patch size (16) because each image is divided by 196 patches, each size is 16x16.
3. If you just want to extract an attention map, I would recommend @rwightman's solution, which is more convenient.

Thank you.

Hankyul

[tomos7231]

Hi @hankyul2

Well understood. Thank you for reply!

[rwightman]

So, this doesn't include the visualization helpers yet, but have added a simpler extraction helper to get the attention activations via one of two methods, fx or hooks.

WIP but can be seen https://github.com/huggingface/pytorch-image-models/pull/2168/files#diff-358e0d5feb2c109ff53d21bc4fa8a6af94566be622b0f1167316216b0036b8b3

import timm
import torch
from timm.utils import AttentionExtract
timm.layers.set_fused_attn(False)
mm = timm.create_model('vit_base_patch16_224')
input = torch.randn(2,3,224,224)
ee = AttentionExtract(mm, method='fx')
oo = ee(input)

for n, t in oo.items():
    print(n, t.shape)
blocks.0.attn.softmax torch.Size([2, 12, 197, 197])
blocks.1.attn.softmax torch.Size([2, 12, 197, 197])
blocks.2.attn.softmax torch.Size([2, 12, 197, 197])
blocks.3.attn.softmax torch.Size([2, 12, 197, 197])
blocks.4.attn.softmax torch.Size([2, 12, 197, 197])
blocks.5.attn.softmax torch.Size([2, 12, 197, 197])
blocks.6.attn.softmax torch.Size([2, 12, 197, 197])
blocks.7.attn.softmax torch.Size([2, 12, 197, 197])
blocks.8.attn.softmax torch.Size([2, 12, 197, 197])
blocks.9.attn.softmax torch.Size([2, 12, 197, 197])
blocks.10.attn.softmax torch.Size([2, 12, 197, 197])
blocks.11.attn.softmax torch.Size([2, 12, 197, 197])

[rwightman]

import timm
import torch
from timm.utils import AttentionExtract
timm.layers.set_fused_attn(False)
mm = timm.create_model('vit_base_patch16_224')
input = torch.randn(2,3,224,224)
ee = AttentionExtract(mm, method='hook')
oo = ee(input)

for n, t in oo.items():
    print(n, t.shape)
blocks.0.attn.attn_drop torch.Size([2, 12, 197, 197])
blocks.1.attn.attn_drop torch.Size([2, 12, 197, 197])
blocks.2.attn.attn_drop torch.Size([2, 12, 197, 197])
blocks.3.attn.attn_drop torch.Size([2, 12, 197, 197])
blocks.4.attn.attn_drop torch.Size([2, 12, 197, 197])
blocks.5.attn.attn_drop torch.Size([2, 12, 197, 197])
blocks.6.attn.attn_drop torch.Size([2, 12, 197, 197])
blocks.7.attn.attn_drop torch.Size([2, 12, 197, 197])
blocks.8.attn.attn_drop torch.Size([2, 12, 197, 197])
blocks.9.attn.attn_drop torch.Size([2, 12, 197, 197])
blocks.10.attn.attn_drop torch.Size([2, 12, 197, 197])
blocks.11.attn.attn_drop torch.Size([2, 12, 197, 197])

[rwightman]

fx uses FX to grab activations by node names, so has high granularity in terms of what can be grabbed as long as you can find the name of the nodes you're interested in. '*attn.softmax' is the default and covers most attention models in timm. But concrete names or other combinations of wildcards can be passed via the names= parameter.

hook method uses old torch hooks, so you need a module to hook. Many attention models in timm have the attn_drop nn.Dropout module right after the attention softmax, so that's useable as pre (or post forward hook if the droput is inactive which is default). ['*attn_drop'] covers most models that have the attention dropout but not all do. softmax was widely implemented in functional form so no module to hook directly.

This extractor can be used to grab other activations too, you could get every norm in a vit ['*.norm'], etc...

[rwightman]

I've tried a number of models and it works pretty well with many without needing different args or custom modifications...

import timm
import timm.utils
import torch
timm.layers.set_fused_attn(False)
mm = timm.create_model('swin_base_patch4_window7_224')
input = torch.randn(2,3,224,224)
ee = timm.utils.AttentionExtract(mm, method='fx')
oo = ee(input)

for n, t in oo.items():    
    print(n, t.shape)
layers.0.blocks.0.attn.softmax torch.Size([128, 4, 49, 49])
layers.0.blocks.1.attn.softmax torch.Size([128, 4, 49, 49])
layers.1.blocks.0.attn.softmax torch.Size([32, 8, 49, 49])
layers.1.blocks.1.attn.softmax torch.Size([32, 8, 49, 49])
layers.2.blocks.0.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.1.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.2.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.3.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.4.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.5.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.6.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.7.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.8.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.9.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.10.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.11.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.12.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.13.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.14.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.15.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.16.attn.softmax torch.Size([8, 16, 49, 49])
layers.2.blocks.17.attn.softmax torch.Size([8, 16, 49, 49])
layers.3.blocks.0.attn.softmax torch.Size([2, 32, 49, 49])
layers.3.blocks.1.attn.softmax torch.Size([2, 32, 49, 49])