Getting nn.MultiHeadAttention attention weights for each head

username: ironcadiz
cooked: <p>I’m using the <code>nn.MultiheadAttention</code> layer (v1.1.0)  with  <code>num_heads=19</code> and an input tensor of size <code>[model_size,batch_size,embed_size]</code></p>
<p>Based on the original Attention is all you need paper, I understand  that there should be a matrix of attention weights for each head (19 in my case), but i can’t find a way of accesing them. When doing a forward pass the returned weights have size <code>[batch_size,model_size,model_size]</code>  instead of something like <code>[batch_size, 19,model_size,model_size]</code>. I’m guessing the weights returned are an average of all the heads but that isn’t specified in the docs.</p>
<p>Is there another way of accessing the full attention weights?</p>

username: vainaijr
cooked: <p>I think we will have to modify F.multi_head_attention_forward for that,<br>
in the end it has,</p>
<pre><code class="lang-auto">if need_weights:
        # average attention weights over heads
        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
        return attn_output, attn_output_weights.sum(dim=1) / num_heads
else:
        return attn_output, None
</code></pre>
<p>if we want full attention weights without averaging over heads, then we will have to change it to,</p>
<pre><code class="lang-auto">if need_weights:
        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
        return attn_output, attn_output_weights
</code></pre>

username: ironcadiz
cooked: <p>I suspected that was going to be the answer <img src="https://discuss.pytorch.org/images/emoji/apple/worried.png?v=9" title=":worried:" class="emoji" alt=":worried:"><br>
Should I make an issue on the repo with the  feature request?</p>
<p>Thanks for your suggestion I will give it a try on a local fork.</p>

username: vainaijr
cooked: <p>There is an active issue on this, to break multi head attention into parts, as it is a bit too long.<br>
I think they will update it soon.<br>
<aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/issues/32590" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/32590" target="_blank" rel="nofollow noopener">[FYI] MultiheadAttention / Transformer</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-01-24" data-time="20:06:25" data-timezone="UTC">08:06PM - 24 Jan 20 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/zhangguanheng66" target="_blank" rel="nofollow noopener">
          <img alt="zhangguanheng66" src="https://avatars0.githubusercontent.com/u/6156351?v=4" class="onebox-avatar-inline" width="20" height="20">
          zhangguanheng66
        </a>
      </div>
    </div>
  </div>
</div>

<div class="github-row">
  <p class="github-content">This issue is created to track the progress to refine nn.MultiheadAttention and nn.Transformer.
Since the release of both modules in PyTorch v1.2.0,...</p>
</div>

<div class="labels">
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">module: nn</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">proposal accepted</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">triaged</span>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>

username: isaacmg
cooked: <p>Is there any progress on this?</p>

username: jbschlosser
cooked: <p>For posterity: a flag to disable averaging of attention weights across heads was added in <a href="https://github.com/pytorch/pytorch/pull/70055" rel="noopener nofollow ugc">#70055</a>. You can now pass <code>average_attn_weights=False</code> to get attention weights per head.</p>