Register_hook not always working (saving gradients)

autograd
1

I’m writing an implementation of the GradCAM algorithm to visualize my neural net. In order to achieve this I’m making use of hooks as so:

...
def save_gradients(self, in):
    self.gradients = in

...

def forward(self, x)
    if i == target_layer:
        x.register_hook(self.save_gradients)
        out = x
...

The problem is this only sometimes works, depending on what layer I’m trying to look at. For example:

Sequential(                                                                                                                                          
  (Conv2d): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)                                                                                                                    
  (BatchNorm2d): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)                                                                                              
  (activation): LeakyReLU(negative_slope=0.1, inplace=True)
)

will work, but not:

Sequential(                                                                                                                                [124/1941]
  (Conv2d): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)                                                                                                   
  (BatchNorm2d): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)                                                                                             
  (activation): LeakyReLU(negative_slope=0.1, inplace=True)                                                                                                                                    
)

Is my intuition/knowledge wrong about something, hence why I can only look at certain layers, or is my code wrong elsewhere?

2

Hi,

What do you mean by it sometimes work?
You should make sure that the hook is actually registered at every forward pass.

3

Hi,

What do you mean by it sometimes work?

When it doesn’t work, self.gradients is None.

You should make sure that the hook is actually registered at every forward pass.

How can I ensure this?

4

If self.gradients is None, this is most likely because the hook was never called because it wasn’t registered.
For debugging purposes, you can add some prints when the register_hook is called for you to check that it happens for all the modules that you want.

5

When self.gradients is None, nothing seems to get called or registered, consequently nor do my prints inside my hook :frowning_face:. forward() just returns and I later get a nice AttributeError: 'NoneType' object has no attribute 'cpu' because self.gradients was None (here I do guided_gradients = self.GradCAM.gradients.cpu().data.numpy()[0])

Are there any subtle edge cases where a hook wouldn’t be properly registered? i.e. if a layer is somehow “un-hookable”?

6

register_hook will always be called.
It won’t only in two cases:

  • It was registered on a Tensor for which the gradients was never computed.
  • The register_hook function is some part of your code that did not run during the forward.
7

For my case, I am implementing gradient clipping as the grad is calculated. I have a print statement before the call to register_hook() and it is indeed being called. The code is as follows and was adapted from @tom’s tutorial.

def forward(self, x):        
    r_out, hidden = self.lstm(x, None)
        
    h = hidden[0]
    if h.requires_grad:
        h.register_hook(lambda x: x.clamp(min=-10, max=10))

Any tips on how to debug this further?

8

This code looks good. What is the issue that you have?

10

:sweat_smile: My apologies, I meant to include a description of the error:

line 101, in <lambda> h.register_hook(lambda x: x.clamp(min=-10, max=10))
AttributeError: 'NoneType' object has no attribute 'clamp'

It seems the tensor being passed to the lambda is None

11

How do you use h later on?
None means that it is all 0s in this context. But I am surprised the hook actually gets called with None…

13

That’s good to know. So an uninitialized tensor defaults to None?

I seem to have a work around by using:

for h in hidden:
    if h.requires_grad:
        h.register_hook(lambda x: x.clamp(min=-10, max=10) if x is not None else x)

Upon examining the contents of the tensors, they contain non-zero values.