martinsbruveris / tensorflow-image-models

TensorFlow port of PyTorch Image Models (timm) - image models with pretrained weights.
https://tfimm.readthedocs.io/en/latest/
Apache License 2.0
286 stars 25 forks source link

Expose attention matrices for ViT architectures #74

Closed Ostaijen closed 1 year ago

Ostaijen commented 1 year ago

Modify the ViT base class to return the attention matrix per block along with the features, under a new key block_{i}_attn.