Extracting Grad from nn.Sequential()

1

Also, I wanted to check if my model may face vanishing gradient issue, therefore I wanted to plot grad thorough out, my model. Please can you help me how can I access them?
Also, any suggestion on the points on which I can look is strongly welcome.
My model,

class basic_CNN2(nn.Module):
    """
    Basic CNN architecture deployed
    """
    def __init__(self):
        super(basic_CNN2, self).__init__()
        self.Block1 = nn.Sequential(
            nn.Conv2d(3,16,3),
            nn.ReLU(),
            nn.Conv2d(16,16,3),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.Block2 = nn.Sequential(
            nn.Conv2d(16,32,3),
            nn.ReLU(),
            nn.Conv2d(32,32,3),
            nn.ReLU(),
            nn.BatchNorm2d(32),
            nn.MaxPool2d(2)
        )
        self.Block3 = nn.Sequential(
            nn.Conv2d(32,64,3),
            nn.ReLU(),
            nn.Conv2d(64,64,3),
            nn.ReLU(),
            nn.BatchNorm2d(64),
            nn.MaxPool2d(2)
        )
        self.Block4 = nn.Sequential(
            nn.Conv2d(64,128,3),
            nn.ReLU(),
            nn.Conv2d(128,128,3),
            nn.ReLU(),
            nn.BatchNorm2d(128),
            nn.MaxPool2d(2),
            nn.Dropout(0.2)
        )
        self.Block5 = nn.Sequential(
            nn.Conv2d(128,256,3),
            nn.ReLU(),
            nn.Conv2d(256,256,3),
            nn.ReLU(),
            nn.BatchNorm2d(256),
            nn.MaxPool2d(2),
            nn.Dropout(0.2)
        )
        self.Linear = nn.Sequential(
            nn.Linear(2304,1024),
            nn.ReLU(),
            nn.Dropout(0.7),
            nn.Linear(1024,512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512,64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64,1),
            nn.Sigmoid()
        )
        

    def forward(self, x):
        x = self.Block1(x)
        x = self.Block2(x)
        x = self.Block3(x)
        x = self.Block4(x)
        x = self.Block5(x)
        x = x.view(-1,self.num_flat_features(x))
        x = self.Linear(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

Thanks in advance

2

Hi,

The simplest is to use t.register_hook(your_fn) on any Tensor for which you want to be able to access the gradients.
You can add there either during the forward on your x Tensors or by adding mod.register_forward_hook(fw_hook) on the modules of your model to get access to these Tensors.

3

I didn’t understand properly, where you trying to me implement in the following manner

for epoch in epochs:
    for idx, batch in enumerate(test_loader):
        x, y = batch['img'], batch['label']
        h = x.register_hook()
        output = model(x)
        optim.zero_grad()
        loss = criterion(output,y)
        loss.backward()
        optimizer.step()
    writer.add_histogram("conv1grad",h[0],epoch)
    writer.add_histogram("conv2grad",h[1],epoch)

Something like this !?
Also can you please tell me how to index inside Sequential Block if my implementation is correct.
Thanks

4

Hi,

x.register_hook(your_fn) will add a hook. And when the gradient for x is computed, your_fn will be called with the gradient that was computed for x.
So then from that function you can do what you want with the gradient (like storing it in a global state to access it later).
You can check the doc for register hook, but what it returns is a handle to be able to remove the hook that’s it.

To access modules in a Sequential, you can index it like any other list: mod[0], mod[1], etc/

5

The problem is I am not able to find a way to get grad of conv1 and conv2 from Sequential Block1 and similarly I wanted to track all way along. Therefore I don’t know what to write in your_func

by mod[0] (mod being a sequential block) will retrieve nn.Conv2d(3,16,3) (in reference to my model var Block1)

Therefore it will be a great help if could help me in an example using my model

Thanks in advance

6

I guess something like:

global_grad_saver = {}

# Have a hook fn that saves into the global dict with a given name
def get_hook(param_name):
    def hook_fn(grad):
        global_grad_saver[param_name] = grad
    return hook_fn


# In your forward code:
    def forward(self, x):
        x = self.Block1(x)
        x.register_hook(get_hook("block1_output"))
        x = self.Block2(x)
        x.register_hook(get_hook("block2_output"))
        # Rest of your forward function


# For a given module:
def get_fw_hook(mod_name):
    def fw_hook(mod, input, output):
        output.register_hook(get_hook(mod_name + "_output"))
    return fw_hook

# Before running the forward of your model
mod.Block1[0].register_forward_hook("Block1.conv0")
mod.Block1[2].register_forward_hook("Block1.conv2")