Obtaining outputs and attention weights from intermediate Transformer layers

username: partially_observed
cooked: <p>Hi,</p>
<p>I am starting to use <code>nn.TransformerEncoder</code> for some experiments and was wondering if there was a way to obtain the outputs and attention weights from intermediate layers?</p>
<p>For example, Table 7 in the <a href="https://arxiv.org/pdf/1810.04805.pdf" rel="nofollow noopener">BERT paper</a>, studies the feature extraction capabilities of BERT and utilizes outputs from intermediate layers. The attention weights would generally help with analyzing the results.</p>
<p>Hence, I was wondering if there was a way to obtain them easily.</p>
<p>Thank you.</p>

username: ypsoh
cooked: <p>I’m not sure what you mean by <code>obtain</code> is, but you can definitely inspect, dump the weights for any layer in the transformer. If you set up for transformer model something like this:<br>
<code>model = Transformer(args)</code>, you can print/access he intermediate weights like this <code>model.encoder.layers[n].self_attn.W_Q</code>.<br>
Of course the exact keys(<code>encoder</code>, <code>self_attn</code>, etc) would change based on how you set up your model but it’s pretty straightforward.</p>

username: EvanZ
cooked: <p>I have a TransformerEncoder model that has been trained on a task similar to BERT’s masking. I would like to access, for example, the outputs of the 2nd to last layer while doing the inference step. Is the idea that I should take the trained model and extract some of the layers up until the last one I want and then do inference?</p>

username: yoshitomo-matsubara
cooked: <p>Hi <a class="mention" href="/u/partially_observed">@partially_observed</a> and <a class="mention" href="/u/evanz">@EvanZ</a></p>
<p>If you’d like to extract input and/or output from some modules in your model, you can easily achieve it by using <code>ForwardHookManager</code> in <a href="https://github.com/yoshitomo-matsubara/torchdistill" rel="noopener nofollow ugc">torchdistill</a>.<br>
Here is <a href="https://github.com/yoshitomo-matsubara/torchdistill/blob/main/demo/extract_intermediate_representations.ipynb" rel="noopener nofollow ugc">an example notebook</a> to demonstrate the feature.</p>

username: EvanZ
cooked: <p>Thanks <a class="mention" href="/u/yoshitomo-matsubara">@yoshitomo-matsubara</a>. I did end up using a hook via <code>register_forward_hook</code> and that is working well.</p>

username: yoshitomo-matsubara
cooked: <p>Great to hear that you resolved it yourself. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
<code>ForwardHookManager</code> in torchdistill does leverage the feature of forward hook in PyTorch like you did. It also takes care of managing <code>handle</code> produced by forward hooks for you, so it should be worth taking a look next time.</p>

username: pacert
cooked: <p>Hey <a class="mention" href="/u/evanz">@EvanZ</a> I want to visualize the attention over a sequence to draw conclusions from it. I use the nn.TransformerEncoder model. The only solution that seems to be available is using forward hooks. I was wondering if you could share how exactly you did this?</p>

username: PaulOppelt
cooked: <p>do you know what the reason could be if W_Q is None?<br>
I initialized my Transformer-Encoder in the following way:<br>
encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout)<br>
transformer = nn.TransformerEncoder(encoder_layer, num_layers)</p>
<p>other than that the network works as it should…</p>

username: joseph-pq
cooked: <p>Hi!<br>
I have this gist file that demonstrate how to visualize attention weights without using an external package but my modifying the TransformerEncoderLayer class:</p><aside class="onebox githubgist" data-onebox-src="https://gist.github.com/joseph-pq/08f4d3757b49b43bb17c67413f962f46">
  <header class="source">

      <a href="https://gist.github.com/joseph-pq/08f4d3757b49b43bb17c67413f962f46" target="_blank" rel="noopener nofollow ugc">gist.github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://gist.github.com/joseph-pq/08f4d3757b49b43bb17c67413f962f46" target="_blank" rel="noopener nofollow ugc">https://gist.github.com/joseph-pq/08f4d3757b49b43bb17c67413f962f46</a></h4>

  <h5>visualize_attention_weights.py</h5>
  <pre><code class="Python">import torch

feature_size = 4  # Example feature size
heads = 4  # Number of attention heads
layers = 1  # Number of transformer layers
dropout = 0.0  # Dropout rate
activation = "relu"  # Activation function
seq_len = 3
batch_size = 2
</code></pre>
   This file has been truncated. <a href="https://gist.github.com/joseph-pq/08f4d3757b49b43bb17c67413f962f46" target="_blank" rel="noopener nofollow ugc">show original</a>

<p>
</p>

  </article>

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

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