How to visualize attention map

[piantic]

Hi,

I want to visualize attention map. I found https://github.com/jeonsworld/ViT-pytorch/blob/main/visualize_attention_map.ipynb

In this repo, I did not found vis option for attention map. (If any, please let me know and I'd appreciate it.)

So, I decided to add this to model.py. like this:

# In VisionTransformer
def forward(self, x):
    feat, attn_weights = self.extract_features(x)

    # classifier
    logits = self.classifier(feat[:, 0])
    return logits, attn_weights

# In Encoder
def forward(self, x):
    attn_weights = []
    out = self.pos_embedding(x)

    for layer in self.encoder_layers:
        out, weights = layer(out)
        attn_weights.append(weights)

    out = self.norm(out)
    return out, attn_weights

# In SelfAttention
def forward(self, x):
    b, n, _ = x.shape

    q = self.query(x, dims=([2], [0]))
    k = self.key(x, dims=([2], [0]))
    v = self.value(x, dims=([2], [0]))

    q = q.permute(0, 2, 1, 3)
    k = k.permute(0, 2, 1, 3)
    v = v.permute(0, 2, 1, 3)

    attn_weights = torch.matmul(q, k.transpose(-2, -1)) / self.scale
    attn_weights = F.softmax(attn_weights, dim=-1)
    out = torch.matmul(attn_weights, v)
    out = out.permute(0, 2, 1, 3)

    out = self.out(out, dims=([2, 3], [0, 1]))

    return out, attn_weights

And I got the result.

But I don't know that it is right or not. Because the result of attention map above link is quite different for me. (I used pretrained weights in here).

I am not sure if my results are correct. I would be happy if I could hear the answer.

Thanks.

[tczhangzhi]

Looks good to me but one thing you should pay attention to is that vit-model-1 is finetuned on the cassava-leaf-disease-classification task. You may expect to visualize an image from that dataset. It is quite different from object classification and focuses on the low-level texture of the input leaf. To visualize the attention map of a dog, you can utilize pre-trained models here.

Anyway, it is a good first try. I'm still hesitating about the operation of extracting the "attention Map" since I don't want it to affect the inference process, that is, to modify the forward function. Maybe later I will check some best practices about hooks. If u r willing to, u can make a PR of your implement.

[piantic]

Thanks for answer. I used your recommended pre-trained models.

Here is result for a dog.

Original attention map in repo for https://github.com/jeonsworld/ViT-pytorch/blob/main/visualize_attention_map.ipynb is below:

It seems to be something, but I'm not sure. What do you think about this part?

If you think this part is okay, It seems that a simple flag-vis can minimize the influence of inference.

If vis is False, Model is working original forward function.

[cifkao]

For people still looking for a solution, my package NoPdb allows capturing attention weights from pretty much any Transformer implementation without any modifications to the code. See a Colab notebook showing how to do this for ViT (a different implementation).

In this case, it would be something like:

with nopdb.capture_calls(SelfAttention.forward) as calls:
    logits = model(x)

calls[0].locals["attn_weights"]  # attention weights of the first layer

[Suryanshg]

Hi, when I try to implement the changes by @piantic, this is the error I am getting:

Traceback (most recent call last): File "C:\Users\Surya\Desktop\Automatic-Pain-Estimation-MQP\scripts\Visualize_Attention_Map.py", line 96, in result_img = get_attention_map(viz_image) File "C:\Users\Surya\Desktop\Automatic-Pain-Estimation-MQP\scripts\Visualize_Attention_Map.py", line 25, in get_attention_map logits, att_mat = model(x.unsqueeze(0)) File "C:\Users\Surya\AppData\Local\Programs\Python\Python39\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl return forward_call(*input, **kwargs) File "C:\Users\Surya\AppData\Local\Programs\Python\Python39\lib\site-packages\vision_transformer_pytorch\model.py", line 269, in forward feat, attn_weights = self.extract_features(x) File "C:\Users\Surya\AppData\Local\Programs\Python\Python39\lib\site-packages\vision_transformer_pytorch\model.py", line 265, in extract_features
feat = self.transformer(emb) File "C:\Users\Surya\AppData\Local\Programs\Python\Python39\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl return forward_call(*input, **kwargs) File "C:\Users\Surya\AppData\Local\Programs\Python\Python39\lib\site-packages\vision_transformer_pytorch\model.py", line 177, in forward out, weights = layer(out) File "C:\Users\Surya\AppData\Local\Programs\Python\Python39\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl return forward_call(*input, **kwargs) File "C:\Users\Surya\AppData\Local\Programs\Python\Python39\lib\site-packages\vision_transformer_pytorch\model.py", line 139, in forward out = self.dropout(out) File "C:\Users\Surya\AppData\Local\Programs\Python\Python39\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl return forward_call(*input, **kwargs) File "C:\Users\Surya\AppData\Local\Programs\Python\Python39\lib\site-packages\torch\nn\modules\dropout.py", line 58, in forward return F.dropout(input, self.p, self.training, self.inplace) File "C:\Users\Surya\AppData\Local\Programs\Python\Python39\lib\site-packages\torch\nn\functional.py", line 1169, in dropout return VF.dropout(input, p, training) if inplace else _VF.dropout(input, p, training) TypeError: dropout(): argument 'input' (position 1) must be Tensor, not tuple

Is there anything else I need to do? I feel that there might be some change that needs to be made in the EncoderBlock part of the model.py file

[piantic]

Hi, @Suryanshg.

This is my example notebook for visualizing attention map using this github. https://www.kaggle.com/code/piantic/vision-transformer-vit-visualize-attention-map/notebook

And you can see visualized version of ViT in below link. https://www.kaggle.com/datasets/piantic/visiontransformerpytorch121

I hope this helps you. Thanks.