Register_full_backward_hook returns None

#1

The following is included in my model’s __init__ function

self.grads = None
target_layer = self.model.features
target_layer.register_backward_hook(self.grad_hook)

And grad_hook(...) is defined as

def grad_hook(self, module, grad_input, grad_output):
        self.grads = grad_output[0].clone().detach()

This code works perfectly and I’m able to retrieve the gradients saved in self.grads after calling .backward(), but it prints the following warning:

UserWarning: 
Using a non-full backward hook when the forward contains multiple autograd Nodes is deprecated and will be removed in future versions. 
This hook will be missing some grad_input. 
Please use register_full_backward_hook to get the documented behavior.
warnings.warn("Using a non-full backward hook when the forward contains multiple autograd Nodes "

However, when I replace register_backward_hook by register_full_backward_hook in __init__, the gradients are not saved, i.e., self.grads is None

Why is this happening?

#2

I cannot reproduce the issue with e.g. AlexNet:

grads = 0
def grad_hook(module, grad_input, grad_output):
    global grads
    grads = grad_output[0].clone().detach()
    

model = models.alexnet()
model.features.register_backward_hook(grad_hook)

x = torch.randn(1, 3, 224, 224)
out = model(x)
out.mean().backward()
print(grads.abs().sum())
> tensor(0.3220)

model = models.alexnet()
model.features.register_full_backward_hook(grad_hook)
grads = 0
out = model(x)
out.mean().backward()
print(grads.abs().sum())
> tensor(0.3470)

Could you post the model definition, so that we could try to reproduce it?

#3

The following warning is printed when I use register_full_backward_hook and it’s not printed when using register_backward_hook.
I’m using the pretrained SqueezeNet model from torchvision (torchvision.models.squeezenet1_1(pretrained = True))

UserWarning: Output 0 of BackwardHookFunctionBackward is a view and is being modified inplace. 
This view was created inside a custom Function (or because an input was returned as-is) 
and the autograd logic to handle view+inplace would override the custom backward associated 
with the custom Function, leading to incorrect gradients. 
This behavior is deprecated and will be forbidden starting version 1.6. 
You can remove this warning by cloning the output of the custom Function. 
(Triggered internally at  /pytorch/torch/csrc/autograd/variable.cpp:547.)
 result = torch.relu_(input)

This stackoverflow question seems to detail the same problem as well.

#4

I still cannot reproduce this issue using my code snippet and by replacing models.alexnet() with models.SqueezeNet(). Both hooks return a valid grad value and the first approach will raise the expected deprecation warning.
Were you able to run into the error using my code snippet?

Note that I’m using a source build with:

torch.__version__
'1.9.0a0+git6d45d7a'

torchvision.__version__
'0.11.0a0+882e11d'
#5

Running your code also reproduces my issue, I’m using versions:

torch: 1.8.1+cu101
torchvision: 0.9.1+cu101

grad_hook is not called when using register_full_backward_hook.

grads = None
def grad_hook(module, grad_input, grad_output):
    print("grad_hook called")
    global grads
    grads = grad_output[0].clone().detach()

# 1
model = models.alexnet()
model.features.register_backward_hook(grad_hook)
x = torch.randn(1, 3, 224, 224)
out = model(x)
out.mean().backward() # prints "grad_hook called"
print(grads.size()) # prints torch.Size([1, 256, 6, 6])
print(grads.abs().sum()) # prints tensor(0.3186)

# 2
model = models.alexnet()
model.features.register_full_backward_hook(grad_hook)
grads = 1234
x = torch.randn(1, 3, 224, 224)
out = model(x)
out.mean().backward()
print(grads) # prints 1234
#6

In that case you might need to update to the latest nightly release or build PyTorch from source, as it seems that a newer version might have fixed this issue.