pytorch hook

register_hook.py

layer_dict = {}

def inspect_layer(module, input, output):
    global layer_dict
    layer_dict[module] = (input, output)
    
handles = [layer.register_forward_hook(inspect_layer) for layer in output_layers]

# handles[0].remove()


# class version (can be a nn.Module)

class Hooks:
    
    def __init__(self, modules):
        self.hooks = [m.register_forward_hook(self.hook_fn) for m in modules]
        self.hooked_values = {}
        
    def hook_fn(self, module, input, output):
        self.hooked_values[module] = (input, output)