microsoft / Swin-Transformer

This is an official implementation for "Swin Transformer: Hierarchical Vision Transformer using Shifted Windows".
https://arxiv.org/abs/2103.14030
MIT License
13.56k stars 2.03k forks source link

Attention Maps visualization #168

Open imanuelroz opened 2 years ago

imanuelroz commented 2 years ago

Hi, I am trying to visualize the self attention maps on my input image trying to extract them from the model, unfortunately I am not succeeding in this task. Could you give me a hint on how to do it? I share you the part of the code in which I am trying to do this task.

attention_maps = []
for module in model.modules():
    if hasattr(module,'attention_patches'): 
        print(module.attention_patches.shape)
        if module.attention_patches.numel() == 224*224:
           attention_maps.append(module.attention_patches)

for attention_map in attention_maps:
    attention_map = attention_map.reshape(224, 224, 1)
    plt.imshow(sample['image'].permute(1, 2, 0), interpolation='nearest')
    plt.imshow(attention_map, alpha=0.7, cmap=plt.cm.Greys)
    plt.show()
AndreiRoibu commented 1 year ago

Hi. I am also trying to do this for my 3D Swin Network. @imanuelroz have you managed to solve the issue? Or has anybody else got any solutions that could help? Thanks!