Unexpected Input gradient tensor shape,, Could help me??

Hi, I tried to visualize for regression task in vaniila_backprop.py. but I faced a problem. Could you help me ?
My problem is input gradient tensor shape mismatch. My hook_layers is as follows.

    def hook_layers(self):
        def hook_function(module, grad_in, grad_out):
            self.gradients = grad_in[0]
            print(grad_in[0].shape, grad_out[0].shape)

        layers = list(self.model._modules['features']) + list(self.model._modules['classifier'])
        print('{}'.format(layers))
        first_layer = layers[0] 
        first_layer.register_backward_hook(hook_function)

My layers output is as follows.

Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)), ReLU(inplace=True), MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False), Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)), ReLU(inplace=True), MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False), Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)), ReLU(inplace=True), MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False), Flatten(), ReLU(inplace=True), Linear(in_features=126720, out_features=1, bias=True)

When I print gradient input[0] and output[0], I expected

torch.Size([1, 3, 180, 360]) torch.Size([1, 32, 180, 360])

But I get

torch.Size([1, 32, 180, 360]) torch.Size([1, 32, 180, 360])

What's wrong in my code..?
My vanilla_backprop.py is as follows.

from Parts import traindataset
from Models import Model 
import torch
from torch.utils.data import Dataset, DataLoader

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class VanillaBackprop():
    def __init__(self, model):
        self.model = model
        self.gradients = None
        self.model.eval()
        self.hook_layers()

    def hook_layers(self):
        def hook_function(module, grad_in, grad_out):
            self.gradients = grad_in[0]
            print(grad_in[0].shape, grad_out[0].shape)

        layers = list(self.model._modules['features']) + list(self.model._modules['classifier'])
        print('{}\n'.format(layers))
        first_layer = layers[0] 
        first_layer.register_backward_hook(hook_function)

    def generate_gradients(self, input_image):
        model_output = self.model(input_image)
        self.model.zero_grad()
        model_output.backward()
        gradients_as_arr = self.gradients.data.cpu().numpy()[0]
        return gradients_as_arr

if __name__ == '__main__':
    model = Model(1).to(device)
    for param_tensor in model.state_dict():
        print(param_tensor, "\t", model.state_dict()[param_tensor].size())
    # model.load_state_dict(torch.load('train/train.h5'))
    trainset = traindataset(1, 1, 3)
    trainloader = DataLoader(trainset, batch_size=1, shuffle=False)
    img_prop, label  = next(iter(trainloader))
    img = img_prop[:, :3, :, :]
    print('input image shape : {}, label shape : {}'.format(img_prop.shape, label.shape))
    # Guided backprop
    VBP = VanillaBackprop(model)
    # Get gradients
    guided_grads = VBP.generate_gradients(img.cuda())

When I printed input image and label shape, I get

input image shape : torch.Size([1, 6, 180, 360]), label shape : torch.Size([1])

My model is as follows.

import torch.nn as nn 
import numpy as np 

class Model(nn.Module):
   def __init__(self, num_answer):
       super(Model, self).__init__()

       self.features = nn.Sequential(
           nn.Conv2d(3, 32, kernel_size=(3, 3), padding=1, stride=1, bias=True),
           nn.ReLU(inplace=True),
           nn.MaxPool2d(2),
           nn.Conv2d(32, 64, kernel_size=(3, 3), padding=1, stride=1, bias=True),
           nn.ReLU(inplace=True),
           nn.MaxPool2d(2),
           nn.Conv2d(64, 128, kernel_size=(3, 3), padding=1, stride=1, bias=True),
           nn.ReLU(inplace=True),
           nn.MaxPool2d(2)

       ) 
       self.classifier = nn.Sequential(
           nn.Flatten(),
           nn.ReLU(inplace=True),
           nn.Linear(126720, num_answer)
       )

   def forward(self, x):
       out = self.features(x)
       out = self.classifier(out)
       return out

Thanks!!

utkuozbulak / pytorch-cnn-visualizations

Unexpected Input gradient tensor shape,, Could help me?? #83