I am trying to put backward hooks in my code and getting the gradient of a specific layer is working. However, I can't seem to save the variable to my dictionary. Does anyone know how to do this?
tcav = {}
def backward_hook(module, grad_input, grad_output):
#print('module:', module)
#print('grad_input:', grad_input)
#print('grad_output:', tuple_of_tensors_to_tensor(grad_output)[0].shape)
return grad_output
for name, layer in model.named_modules():
if name.startswith(layer_name):
print("Hook made for...", name)
layer.register_full_backward_hook(backward_hook)
model.eval()
tcav = {}
for ind, (img, label) in enumerate(loader):
img.requires_grad=True
img = img.to(device).float()
output = model(img)
output.mean().backward()
tcav[ind] = output.grad
return tcav
Here, output.grad seems to be None for all images.