`register_forward_hook` doesn't work for NestedTensor

jasperhyp: <p>Hi!  I observed an interesting behavior of <code>register_forward_hook</code> but am not sure if it is indeed the case, so I am posting it here in the hope that someone will help clarify this to me.</p>
<p>First, I initialized a model and loaded a checkpoint from a previous run for inference:</p>
<pre><code class="lang-python">best_model = my_model(args)
checkpoint = torch.load(args.model_save_dir+'/best_model.pt')
best_model.load_state_dict(checkpoint)
best_model = best_model.cpu()
</code></pre>
<p>Then, I tried to make predictions by passing the model to another method:</p>
<pre><code class="lang-python">get_all_preds(best_model, test_batch, other_args)
</code></pre>
<p>This <code>get_all_preds</code> function contains a <code>get_activation</code> function and a dictionary prepared to receive attention weights from the last attention layer:</p>
<pre><code class="lang-python">activation = {}
def get_activation(name, activation):
    def hook(model, input, output):
        activation[name] = output
    return hook

best_model.transformer.transformer_encoder.layers[-1].self_attn.register_forward_hook(get_activation('last-layer-attention', activation))

best_model.eval()
preds = best_model(test_batch)

&lt;Breakpoint&gt;
</code></pre>
<p>I thought I would get an <code>activation</code> with some attention weights, but the <code>activation</code> is empty at the breakpoint.</p>
<p>HOWEVER, if I simply move the code snippet inside <code>get_all_preds</code> to the same method appended after the place <code>best_model</code> where is instantiated, I can get attention weights without problem!</p>
<p>I.e.:</p>
<pre><code class="lang-python">best_model = my_model(args)
checkpoint = torch.load(args.model_save_dir+'/best_model.pt')
best_model.load_state_dict(checkpoint)
best_model = best_model.cpu()

activation = {}
def get_activation(name, activation):
    def hook(model, input, output):
        activation[name] = output
    return hook

best_model.transformer.transformer_encoder.layers[-1].self_attn.register_forward_hook(get_activation('last-layer-attention', activation))

best_model.eval()
preds = best_model(test_batch)

&lt;Breakpoint&gt;
</code></pre>
<p>My question is, why passing the model to another method broke the hook’s pipeline?</p>

JuanFMontesinos: <p>The reason why pytorch uses  <code>forward</code> instad of the python’s default <code>__call__</code> is that all this machinery is coded there.</p>
<p>When you call model(…) you are calling <code>model.__call__(...)</code> which internally calls the forward func.<br>
The hooks are triggered in the <code>__call__</code> func.<br>
If you call forward externally you are skipping this.</p>

jasperhyp: <p>Thanks for your response!  I think I am calling <code>model()</code> in the same way, so why isn’t <code>__call__</code> triggered when calling <code>model(...)</code> inside another method, when the <code>model</code> is declared in the main method?</p>

JuanFMontesinos: <p>It doesn’t make sense <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"><br>
Could you make a reproducible script?<br>
I would say you have some variable scope problem or some silly problem.</p>

jasperhyp: <p>Through a controlled experiment, I think the problem is either in my model or in <code>model.eval()</code> + <code>torch.no_grad</code>.</p>
<p>When <code>@torch.no_grad()</code> is present-- When there’s no <code>model.eval()</code> inside the <code>get_test</code> method, it’s fine.  When it exists, no activation recorded.  Of course, the outputs of these two cases differ a lot because of the Dropout and Normalization layers in the model.</p>
<p>When <code>@torch.no_grad()</code> is not present-- Activation is always recorded.</p>
<p>Although indeed, this doesn’t make sense to me.</p>
<p>The following snippet shows a controlled experiment:</p>
<pre><code class="lang-python">@torch.no_grad()
def get_test(model, test_batch):
    # model.eval()  # Comment and uncomment this
    return model(test_batch)

def get_activation(name, activation):
    def hook(model, input, output):
        activation[name] = output
    return hook

activation = {}
model = MyModel(args)
checkpoint = torch.load(some_path)
model.load_state_dict(checkpoint)
model = model.to(device)

model.transformer.transformer_encoder.layers[-1].self_attn.register_forward_hook(get_activation('last-layer-attention', activation))

print(get_test(model, all_batch))
print(activation)
</code></pre>

JuanFMontesinos: <p>To discard it’s a pytorch issue and not something related to the code of your model you could try to catch an activation for a toy model.</p>
<p><code>__call__</code> implementation</p>
<pre><code class="lang-auto">    def _call_impl(self, *input, **kwargs):
        forward_call = (self._slow_forward if torch._C._get_tracing_state() else self.forward)
        # If we don't have any hooks, we want to skip the rest of the logic in
        # this function, and just call forward.
        if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
                or _global_forward_hooks or _global_forward_pre_hooks):
            return forward_call(*input, **kwargs)
        # Do not call functions when jit is used
        full_backward_hooks, non_full_backward_hooks = [], []
        if self._backward_hooks or _global_backward_hooks:
            full_backward_hooks, non_full_backward_hooks = self._get_backward_hooks()
        if _global_forward_pre_hooks or self._forward_pre_hooks:
            for hook in (*_global_forward_pre_hooks.values(), *self._forward_pre_hooks.values()):
                result = hook(self, input)
                if result is not None:
                    if not isinstance(result, tuple):
                        result = (result,)
                    input = result

        bw_hook = None
        if full_backward_hooks:
            bw_hook = hooks.BackwardHook(self, full_backward_hooks)
            input = bw_hook.setup_input_hook(input)

        result = forward_call(*input, **kwargs)
        if _global_forward_hooks or self._forward_hooks:
            for hook in (*_global_forward_hooks.values(), *self._forward_hooks.values()):
                hook_result = hook(self, input, result)
                if hook_result is not None:
                    result = hook_result

        if bw_hook:
            result = bw_hook.setup_output_hook(result)

        # Handle the non-full backward hooks
        if non_full_backward_hooks:
            var = result
            while not isinstance(var, torch.Tensor):
                if isinstance(var, dict):
                    var = next((v for v in var.values() if isinstance(v, torch.Tensor)))
                else:
                    var = var[0]
            grad_fn = var.grad_fn
            if grad_fn is not None:
                for hook in non_full_backward_hooks:
                    wrapper = functools.partial(hook, self)
                    functools.update_wrapper(wrapper, hook)
                    grad_fn.register_hook(wrapper)
                self._maybe_warn_non_full_backward_hook(input, result, grad_fn)

        return result
</code></pre>
<p>Apparently there is no condition to skip the hooking system if no grad or eval mode. I’d say it should work in both cases.</p>

jasperhyp: <p>I agree.  It also has nothing to do with loading checkpoint / lazy initialization.  Now the only problem seems to be the model itself…  That’s a headache.</p>

jasperhyp: <p>Ahhhhhh SHIT I got it!  It’s because of the <em>fast path</em> implemented by <code>torch.nn.TransformerEncoderLayer</code> <a href="https://pytorch.org/docs/stable/generated/torch.nn.TransformerEncoderLayer.html" rel="noopener nofollow ugc">(see here)</a>.  It looks like hook just doesn’t apply to <code>torch.NestedTensor</code>.  When I change any of the premises (leading to a <em>normal path</em>), the hook works as usual!</p>

jasperhyp: <p>Does anyone happen to know how to hook <code>torch.NestedTensor</code>?  It almost took me 6 hours to figure out it is this <em>fast path</em> issue…</p>

JuanFMontesinos: <p>Oh it was that <img src="https://discuss.pytorch.org/images/emoji/apple/expressionless.png?v=12" title=":expressionless:" class="emoji" alt=":expressionless:" loading="lazy" width="20" height="20"><br>
Btw it would be nice if you open a PR about it to give a warning or document it so that people is aware.</p>