Pytorch Transformer doesn't work with register_forward_pre_hook

username: thyeros
cooked: <p>I have a testing code as follows, and it only prints msg from hook_fn registered with Linear.</p>
<p>Is there a reason why hook doesn’t work with _LinearWithBias in MultiheadAttention?</p>
<hr>
<p>import torch<br>
import torch.nn as nn<br>
import torch.nn.functional as F<br>
from torch.nn import TransformerEncoder, TransformerEncoderLayer</p>
<p>class SimpleVT(nn.Module):</p>
<pre><code>def __init__(self):
    super().__init__()
    self.enc_layers = TransformerEncoderLayer(40, 2, 20, 0.5)
    self.encoder = TransformerEncoder(self.enc_layers, 2)
    self.decoder = nn.Linear(40, 2)
    
def forward(self, x):
    x = self.enc_layers(x)
    x = self.encoder(x)
    x = self.decoder(x)
    return x
</code></pre>
<p>model = SimpleVT().cuda()<br>
print(model)</p>
<p>def forward_pre_hook_fn(module, i):<br>
print(“forward_pre_hook_fn”, type(module))</p>
<p>print(type(model.enc_layers.self_attn.out_proj))<br>
model.enc_layers.self_attn.out_proj.register_forward_pre_hook(forward_pre_hook_fn)<br>
model.enc_layers.linear1.register_forward_pre_hook(forward_pre_hook_fn)</p>
<p>x = torch.rand([64, 102, 40]).cuda()<br>
y_hat = torch.rand([64, 102, 2]).cuda()</p>
<p>y = model(x)</p>
<hr>
<p>There is no print from the hook registered with model.enc_layers.self_attn.out_proj.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/9/a9b64b380c4d71033d4f044e7d50aa14f69c7abd.png" data-download-href="https://discuss.pytorch.org/uploads/default/a9b64b380c4d71033d4f044e7d50aa14f69c7abd" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/9/a9b64b380c4d71033d4f044e7d50aa14f69c7abd.png" alt="image" data-base62-sha1="odlfP4Fr3IIbEetqbWPtNaq1kGF" width="458" height="500" data-dominant-color="EBEBEB"><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">718×783 41.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>

username: ptrblck
cooked: <p>I get a valid output:</p>
<pre><code class="lang-python">y = model(x)
forward_pre_hook_fn &lt;class 'torch.nn.modules.linear.Linear'&gt;
</code></pre>
<p>Also, based on the screenshot it seems you are seeing the same output.</p>
<p>PS: you can post code snippets by wrapping them into three backticks ```, which makes debugging easier. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

username: thyeros
cooked: <p>Hi, <a class="mention" href="/u/ptrblck">@ptrblck</a>, Actually, it is not correct. I have registered the same hook_fn on two modules.</p>
<p>model.enc_layers.self_attn.out_proj.register_forward_pre_hook(forward_pre_hook_fn)<br>
model.enc_layers.linear1.register_forward_pre_hook(forward_pre_hook_fn)</p>
<p>What we are seeing is from the 2nd call, and the 1st one doesn’t work: we should see another print. Hope this makes clear.</p>

username: ptrblck
cooked: <p>Thanks for the clarification. I missed this output.<br>
The layer itself is never called, but its parameters are used in <code>F.multi_head_attention_forward</code> in <a href="https://github.com/pytorch/pytorch/blob/b5dd6e3e61c59e88ce6e8aee2fa5d08d6c56fd97/torch/nn/modules/activation.py#L966-L985">these lines on code</a>, which is why the hook isn’t called.</p>

username: thyeros
cooked: <p>oh, I see from the implementation, thanks. For this particular case, I wonder why it is done that way, but I guess it was the design choice. Is there a way to detect this case (there is a child module, but never called in forward) systematically? Or should it be a case-by-case thing?</p>

username: ptrblck
cooked: <p>I don’t know why this approach was chosen and I don’t know if there is another way of checking, if the <code>forward</code> method was called besides what you’ve already did: using hooks and checking their output.</p>