How to use '_register_hooks'?

[yw155]

Hi @pierrestock, I would like to ask you how to use the function '_register_hooks' in watcher.py. Could you provide some example codes? In addition, I notice that It may generate errors during saving a model as torch.save may not support the function 'fwd_hook'. Thank you.

[pierrestock]

Hi yw155,

Thanks for reaching out! The documentation about forward hooks for nn.Module can be found here.

Regarding the function _register_hooks you are referring to, it successively does the following:

• Define a forward hook to save the input activations of each layer;
• Register the forward hook the each layer we are interested in;
• Store a handle for each hook that can be used to remove the hook from the layers it is attached to.

An example of utilisation of this function is given here.

Hope this helps,

Pierre

[yw155]

Thank you for your detailed illustration.