I’m working on CNN saliency maps.
def saliency_map_general(model, input, label):
    if (isinstance(model, torchvision.models.Inception3)):
        input = preprocess_inception(input)
    else:
        input = preprocess(input)

    output = model.forward(input)
    model.zero_grad()
                                       
    output[0][label].backward()

    grads = input.grad.data.clamp(min=0)
    grads.squeeze_()
    grads.transpose_(0,1)
    grads.transpose_(1,2)
    grads = np.amax(grads.cpu().numpy(), axis=2)
    
    true_image = input.data
    true_image = true_image.squeeze()
    true_image = true_image.transpose(0,1)
    true_image = true_image.transpose(1,2)
    true_image = deprocess(true_image)

    fig = plt.figure()
    plt.rcParams["figure.figsize"] = (20, 20)


    a = fig.add_subplot(1,2,1)
    imgplot = plt.imshow(true_image)
    plt.title('Original Image')
    plt.axis('off') 

    a = fig.add_subplot(1,2,2)
    imgplot = plt.imshow(grads, cmap='hot')
    plt.axis('off') 
    plt.title('Saliency Map')
    
    return grads

Here, I’m finding the derivative of the output “true class” w.r.t. input pixels.
Now, I want to use ‘guided backpropagation’ i.e. taking max(grad, 0) at each layer before passing the gradient to previous layer during backward pass. The documentation of register_backward_hook is not very detailed.
I want something like
def hookfunc(model, gi, go):
    grads = gi > 0
    return grads

h = model.register_backward_hook(hookfunc)

I tried doing this, but this hook callback function never runs. I think that is because I’m running .backward() on an element of the output Variable, and not on the model. I’m not sure about this, but the documentation isn’t making anything clear.
To clarify, I want the hook callback function to run for every layer during back pass. Any help on how to do this would be appreciated.