Extracting Grad from nn.Sequential()

username: Vatsal_Malaviya
cooked: <p>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?<br>
Also, any suggestion on the points on which I can look is strongly welcome.<br>
My model,</p>
<pre><code class="lang-auto">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
</code></pre>
<p>Thanks in advance</p>

username: albanD
cooked: <p>Hi,</p>
<p>The simplest is to use <code>t.register_hook(your_fn)</code> on any Tensor for which you want to be able to access the gradients.<br>
You can add there either during the forward on your <code>x</code> Tensors or by adding <code>mod.register_forward_hook(fw_hook)</code> on the modules of your model to get access to these Tensors.</p>

username: Vatsal_Malaviya
cooked: <p>I didn’t understand properly, where you trying to me implement in the following manner</p>
<pre><code class="lang-auto">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)
</code></pre>
<p>Something like this !?<br>
Also can you please tell me how to index inside Sequential Block if my implementation is correct.<br>
Thanks</p>

username: albanD
cooked: <p>Hi,</p>
<p><code>x.register_hook(your_fn)</code> will add a hook. And when the gradient for <code>x</code> is computed, <code>your_fn</code> will be called with the gradient that was computed for <code>x</code>.<br>
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).<br>
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.</p>
<p>To access modules in a Sequential, you can index it like any other list: <code>mod[0]</code>, <code>mod[1]</code>, etc/</p>

username: Vatsal_Malaviya
cooked: <p>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</p>
<p>by mod[0] (mod being a sequential block) will retrieve nn.Conv2d(3,16,3) (in reference to my model var Block1)</p>
<p>Therefore it will be a great help if could help me in an example using my model</p>
<p>Thanks in advance</p>

username: albanD
cooked: <p>I guess something like:</p>
<pre><code class="lang-python">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")
</code></pre>