jeonsworld / ViT-pytorch

Pytorch reimplementation of the Vision Transformer (An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale)
MIT License
1.95k stars 374 forks source link

Get Attention weights #24

Closed joellehanna closed 3 years ago

joellehanna commented 3 years ago

Hi, is there an easy way to get/extract the attention weights in order to visualise the attention map ?

Thanks !