Register forward hook is not working

Ragnor · April 19, 2021, 5:10am

def extract_features(self, t_model, s_model, t_model_cfg, s_model_cfg):
    global t_FM
    t_FM = []
    global s_FM
    s_FM = []


    def hook1(module1, input, output):
        print("hooker Working")
        t_FM.append(output)

    handle = t_model.tbackbone.layer2[1].conv1.register_forward_hook(hook1)
    handle.remove()

    def hook2(module2, input, output):
        print("hooker Working")
        s_FM.append(output)

    handle = s_model.sbackbone.layer2[1].conv1.register_forward_hook(hook2)
    handle.remove()

    return t_FM, s_FM

sio277 · April 19, 2021, 5:18am

Remove the hook handle after fire it (after forwarding).

Ragnor · April 19, 2021, 5:23am

Hello @sio277 ,
can you explain me more.
thank you.

sio277 · April 19, 2021, 5:30am

The registered hooks were removed before they are called. Try something like this:

    t_FM = []
    s_FM = []
    hooks = []

    def hook1(module1, input, output):
        print("hooker Working")
        t_FM.append(output)

    hooks.append(t_model.tbackbone.layer2[1].conv1.register_forward_hook(hook1))

    def hook2(module2, input, output):
        print("hooker Working")
        s_FM.append(output)

    hooks.append(s_model.sbackbone.layer2[1].conv1.register_forward_hook(hook2))

    your_model(input_tensors)  # forward with firing all the hooks registered
    for hook_handle in hooks:
        hook_handle.remove()

Also, you don’t need to set t_FM and s_FM as global (nonlocal) variables .

sio277 · April 19, 2021, 5:38am

github.com

pytorch/pytorch/blob/ffdecc1ac48baf6aea3848e9f530256342a7046b/torch/nn/modules/module.py#L1009


            tracing_state.push_scope(name)
        else:
            recording_scopes = False
    try:
        result = self.forward(*input, **kwargs)
    finally:
        if recording_scopes:
            tracing_state.pop_scope()
    return result

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()

Check the source code of nn.Module.__call__(), where all the hooks are fired with forward().

Ragnor · April 19, 2021, 5:57am

Thank you so much @sio277
I global them because of i want to use them in another def