Gradients don't flow back during guided backprop

kbtorcher: <p>Following is my custom NN:</p>
<blockquote>
<p>CNN(<br>
(seq_model): Sequential(<br>
(Conv2d_1): Conv2d(3, 64, kernel_size=(5, 5), stride=(1, 1), padding=same)<br>
(BN_1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)<br>
(relu_1): ReLU()<br>
(Maxpool2d_1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)<br>
(Conv2d_2): Conv2d(64, 128, kernel_size=(5, 5), stride=(1, 1), padding=same)<br>
(BN_2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)<br>
(relu_2): ReLU()<br>
(Maxpool2d_2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)<br>
(Conv2d_3): Conv2d(128, 256, kernel_size=(5, 5), stride=(1, 1), padding=same)<br>
(BN_3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)<br>
(relu_3): ReLU()<br>
(Maxpool2d_3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)<br>
(Conv2d_4): Conv2d(256, 512, kernel_size=(5, 5), stride=(1, 1), padding=same)<br>
(BN_4): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)<br>
(relu_4): ReLU()<br>
(Maxpool2d_4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)<br>
(Conv2d_5): Conv2d(512, 1024, kernel_size=(5, 5), stride=(1, 1), padding=same)<br>
(BN_5): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)<br>
(relu_5): ReLU()<br>
(Maxpool2d_5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)<br>
(Flatten): Flatten(start_dim=1, end_dim=-1)<br>
(fc1): Linear(in_features=16384, out_features=128, bias=True)<br>
(Dropout): Dropout(p=0, inplace=False)<br>
(fcReLU): ReLU()<br>
(head): Linear(in_features=128, out_features=10, bias=True)<br>
)<br>
)</p>
</blockquote>
<p>I am experimenting with guided backprop and this is my code:</p>
<pre><code class="lang-auto">class Guided_backprop:
    def __init__(self, model, utils_agent):
        self.model = model
        self.image_reconstruction = None # store R0
        self.activation_maps = []  # store f1, f2, ...
        for _, p in self.model.named_parameters():
            p.requires_grad = True
        self.model.eval()
        self.register_hooks()

    def register_hooks(self):
        def first_layer_hook_fn(module, grad_in, grad_out):
            self.image_reconstruction = grad_in[0] 

        def forward_hook_fn(module, input, output):
            self.activation_maps.append(output)

        def backward_hook_fn(module, grad_in, grad_out):
            grad = self.activation_maps.pop() 
            # for the forward pass, after the ReLU operation, 
            # if the output value is positive, we set the value to 1,
            # and if the output value is negative, we set it to 0.
            grad[grad &gt; 0] = 1 
            
            # grad_out[0] stores the gradients for each feature map,
            # and we only retain the positive gradients
            new_grad_in = grad * torch.clamp(grad_out[0], min=0.0)
            return (new_grad_in,)

        modules = []
        for module in self.model.seq_model.named_children():
            modules.append(module)

        # travese the modules，register forward hook &amp; backward hook
        # for the ReLU
        for name, module in modules:
            if isinstance(module, nn.ReLU):
                module.register_forward_hook(forward_hook_fn)
                module.register_backward_hook(backward_hook_fn)

        # register backward hook for the first conv layer
        first_layer = modules[0][1]
        first_layer.register_backward_hook(first_layer_hook_fn)

    def visualize(self, datapoint):
        def normalize(image):
            norm = (image - image.mean())/image.std()
            norm = norm * 0.1
            norm = norm + 0.5
            norm = norm.clip(0, 1)
            return norm

        input_image, _ = datapoint
        target_class = None
        input_image = input_image.unsqueeze(0).requires_grad_().to(device)
        model_output = self.model(input_image)
        self.model.zero_grad()
        pred_class = model_output.argmax().item()
        
        grad_target_map = torch.zeros(model_output.shape, dtype=torch.float, device=device)
        
        if target_class is not None:
            grad_target_map[0][target_class.argmax(0).item()] = 1
        else:
            grad_target_map[0][pred_class] = 1
        
        model_output.backward(gradient=grad_target_map)
        input_image = input_image.squeeze(0)
        result = self.image_reconstruction.data[0].permute(1,2,0)
        print("Img reconst", self.image_reconstruction.shape)
        result = normalize(result)
        gbp_result = wandb.Image(result.cpu().numpy(), caption='Guided BP Image')
        orig_img = wandb.Image(self.utils_agent.invTransf(input_image).cpu(), caption='Original Image')

        wandb.log({'Orig_img': orig_img, 'GBP_Result': gbp_result})
</code></pre>
<p><strong>My image reconstruction is of the size</strong>:</p>
<blockquote>
<p>torch.Size([1, 64, 128, 128])</p>
</blockquote>
<p>which is the same as my <strong>conv1</strong> layer output. This means gradients aren’t flowing back through the input layer. My input image size is:  torch.Size([1, 3, 128, 128]). Where am I going wrong?</p>