PyTorch forward hook on dropout layer of vision transformer behaving oddly

username: Display8176
cooked: <p>I want to visualize the attention maps on the vision transformer. As a sample code, I want to register a hook after the self-attention block(line no 100 below). The code from the time repo is produced in the screenshot below. Since the dropout layer functions as the <code>nn.Identity</code> layer in eval mode, the <code>self.attn_drop</code> output or input should be the same, and a hook placed on the dropout layer <code>self.attn_drop</code> should serve the purpose. Here is the minimum working example to test.</p><aside class="onebox githubblob" data-onebox-src="https://github.com/huggingface/pytorch-image-models/blob/6b9e38b5891e1abbf63827239113ec30cbe631aa/timm/models/vision_transformer.py#L100">
  <header class="source">

      <a href="https://github.com/huggingface/pytorch-image-models/blob/6b9e38b5891e1abbf63827239113ec30cbe631aa/timm/models/vision_transformer.py#L100" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/huggingface/pytorch-image-models/blob/6b9e38b5891e1abbf63827239113ec30cbe631aa/timm/models/vision_transformer.py#L100" target="_blank" rel="noopener nofollow ugc">huggingface/pytorch-image-models/blob/6b9e38b5891e1abbf63827239113ec30cbe631aa/timm/models/vision_transformer.py#L100</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="90" style="counter-reset: li-counter 89 ;">
          <li></li>
          <li>        if self.fused_attn:</li>
          <li>            x = F.scaled_dot_product_attention(</li>
          <li>                q, k, v,</li>
          <li>                dropout_p=self.attn_drop.p if self.training else 0.,</li>
          <li>            )</li>
          <li>        else:</li>
          <li>            q = q * self.scale</li>
          <li>            attn = q @ k.transpose(-2, -1)</li>
          <li>            attn = attn.softmax(dim=-1)</li>
          <li class="selected">            attn = self.attn_drop(attn)</li>
          <li>            x = attn @ v</li>
          <li></li>
          <li>        x = x.transpose(1, 2).reshape(B, N, C)</li>
          <li>        x = self.proj(x)</li>
          <li>        x = self.proj_drop(x)</li>
          <li>        return x</li>
          <li></li>
          <li></li>
          <li>class LayerScale(nn.Module):</li>
          <li>    def __init__(</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<pre><code class="lang-auto">import timm
import torch

# Prepare Model
vit_model = timm.create_model('vit_large_patch16_384', pretrained=True).to("cpu")
vit_model.eval(); 
x = torch.randn(3, 384, 384)
x.shape

attentions = []
def get_attention(module, input, output):
        attentions.append(output.detach().clone())
        attentions.append(input[0].detach().clone())

model = vit_model.eval().blocks[0].attn.attn_drop
print(model)

hook = model.register_forward_hook(get_attention)
print(model._forward_hooks)

with torch.no_grad():
    output = model.eval()(x.unsqueeze(0))

print(attentions[0].shape,attentions[1].shape)
#(torch.Size([1, 3, 384, 384]), torch.Size([1, 3, 384, 384]))
hook.remove()
attentions = []

print(model._forward_hooks)
</code></pre>
<p>This way the hook was getting registered properly. However, the input and output shape of the dropout layer are the same but incorrect. The correct shape should be <code>(batch_size, num_head, num_token, num_token)</code>, i.e., (1, 16, 577, 577). However, the hook function gives the shape of  [1, 3, 384, 384]. What am I missing?</p>
<p>Also, if I register the hook using a for loop for all the sub-modules of the attention block, then it is possible to register a hook on the dropout layer, but still, the list is empty, and therefore the “list index is out of range.”<br>
Are these two things related? Does the dropout layer have some specific behavior?</p>
<pre><code class="lang-auto">import timm
import torch

# Prepare Model
vit_model = timm.create_model('vit_large_patch16_384', pretrained=True).to("cpu")
vit_model.eval(); 
x = torch.randn(3, 384, 384)
x.shape

class VITAttention:
    def __init__(self, model, attention_layer_name='attn_drop'):
        self.model = model
        self.attentions = []
        self.hook = []

        for name, module in self.model.named_modules():
            if attention_layer_name in name:
                self.hook.append(module.register_forward_hook(self.get_attention))

    def get_attention(self, module, input, output):
        self.attentions.append(output.detach().clone())

    def __enter__(self, *args): return self
    def __exit__(self, *args):
        for handle in self.hook:
            handle.remove()

with VITAttention(vit_model) as hook:
    with torch.no_grad():
        output = hook.model.eval()(x.unsqueeze(0))

    print("Shape of qkv attention " + str(hook.attentions[0].shape))
    print("Shape of key attention " + str(hook.attention[1].shape))
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/b/7b4a34bfd2e0b50f92e3904a8873d5cdd621d356.png" data-download-href="https://discuss.pytorch.org/uploads/default/7b4a34bfd2e0b50f92e3904a8873d5cdd621d356" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/b/7b4a34bfd2e0b50f92e3904a8873d5cdd621d356.png" alt="image" data-base62-sha1="hAFIiUO8pZPHQcW60zhk4T9mnB4" width="690" height="162" data-dominant-color="F3F3F8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">807×190 7.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>