How to register forward hooks for each module

1

I’d like to register forward hooks for each module in my network. I have a working code for one module. The most important part looks this way:

    def __init__(self, model, layer_name):
        self.hook = model.module.linear1.register_forward_hook(self.hook_fn)

I don’t know, however, how to access modules using layer_name variable.

So I suspect there’s some better way to do this. First, I’d like to list all the layers names (model.children() or model.named_parameters() don’t do this) and second, I’d like to use them somehow with the attached code.

2

You could use model.named_modules() to iterate all submodules including their names and store the hooks in a dict:

self.hooks = {}
for name, module in model.named_modules():
    hooks[name] = module.register_forward_hook(self,hook_fn)
7 Likes
3

Cool! I overlooked .named_modules() but I knew there has to be something like that.

4

Hello ptrblck!
I have an additional question…
I want to save inputs (or feature space) of each modules in a dictionary wise
so… for example if I have module names as model.layer1.conv1, model.layer1.conv2
I would like to save their input in a form of
input['model.layer1.conv1'] = [input]
and append them as iteration goes
input['model.layer1.conv1'].append(input)
I would like to do it with forward hook
I’m not sure if the code you presented here does the job for me

thanks a lot

5

.named_modules() will return the name and module recursively for the complete model.
If you want to just register specific modules, such as nn.Conv2d but not e.g. BottleNeck or other custom modules, you could use a condition to check for the module via if isinstance().

6

thanks!
hm… so named_modules will only give me a specific name of the parameter only when I have it iterated over the module…

is there any way to get a specific id of the module without iterating the whole module??
thanks!

7

Yes, you could either directly access and index the module via:

print(model.layer1.conv1)

or use getattr(model, name).