Register_hook not always working (saving gradients)

username: dib
cooked: <p>I’m writing an implementation of the GradCAM algorithm to visualize my neural net. In order to achieve this I’m making use of hooks as so:</p>
<pre><code class="lang-auto">...
def save_gradients(self, in):
    self.gradients = in

...

def forward(self, x)
    if i == target_layer:
        x.register_hook(self.save_gradients)
        out = x
...
</code></pre>
<p>The problem is this only sometimes works, depending on what layer I’m trying to look at. For example:</p>
<pre><code class="lang-auto">Sequential(                                                                                                                                          
  (Conv2d): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)                                                                                                                    
  (BatchNorm2d): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)                                                                                              
  (activation): LeakyReLU(negative_slope=0.1, inplace=True)
)
</code></pre>
<p>will work, but not:</p>
<pre><code class="lang-auto">Sequential(                                                                                                                                [124/1941]
  (Conv2d): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)                                                                                                   
  (BatchNorm2d): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)                                                                                             
  (activation): LeakyReLU(negative_slope=0.1, inplace=True)                                                                                                                                    
)
</code></pre>
<p>Is my intuition/knowledge wrong about something, hence why I can only look at certain layers, or is my code wrong elsewhere?</p>

username: albanD
cooked: <p>Hi,</p>
<p>What do you mean by it sometimes work?<br>
You should make sure that the hook is actually registered at <strong>every</strong> forward pass.</p>

username: dib
cooked: <p>Hi,</p>
<blockquote>
<p>What do you mean by it sometimes work?</p>
</blockquote>
<p>When it doesn’t work, <code>self.gradients</code> is <code>None</code>.</p>
<blockquote>
<p>You should make sure that the hook is actually registered at <strong>every</strong> forward pass.</p>
</blockquote>
<p>How can I ensure this?</p>

username: albanD
cooked: <p>If self.gradients is None, this is most likely because the hook was never called because it wasn’t registered.<br>
For debugging purposes, you can add some prints when the register_hook is called for you to check that it happens for all the modules that you want.</p>

username: dib
cooked: <p>When <code>self.gradients</code> is None, nothing seems to get called or registered, consequently nor do my prints inside my hook <img src="https://discuss.pytorch.org/images/emoji/apple/frowning_face.png?v=9" title=":frowning_face:" class="emoji" alt=":frowning_face:">. <code>forward()</code> just returns and I later get a nice <code>AttributeError: 'NoneType' object has no attribute 'cpu'</code> because <code>self.gradients</code> was None (here I do <code>guided_gradients = self.GradCAM.gradients.cpu().data.numpy()[0]</code>)</p>
<p>Are there any subtle edge cases where a hook wouldn’t be properly registered? i.e. if a layer is somehow “un-hookable”?</p>

username: albanD
cooked: <p><code>register_hook</code> will always be called.<br>
It won’t only in two cases:</p>
<ul>
<li>It was registered on a Tensor for which the gradients was never computed.</li>
<li>The register_hook function is some part of your code that did not run during the forward.</li>
</ul>

username: buchholzmd
cooked: <p>For my case, I am implementing gradient clipping as the grad is calculated. I have a print statement before the call to <code>register_hook()</code> and it is indeed being called. The code is as follows and was adapted from <a class="mention" href="/u/tom">@tom</a>’s  <a href="https://github.com/t-vi/pytorch-tvmisc/blob/master/misc/graves_handwriting_generation.ipynb" rel="noopener nofollow ugc">tutorial</a>.</p>
<pre><code class="lang-auto">def forward(self, x):        
    r_out, hidden = self.lstm(x, None)
        
    h = hidden[0]
    if h.requires_grad:
        h.register_hook(lambda x: x.clamp(min=-10, max=10))
</code></pre>
<p>Any tips on how to debug this further?</p>

username: albanD
cooked: <p>This code looks good. What is the issue that you have?</p>

username: buchholzmd
cooked: <p><img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=9" title=":sweat_smile:" class="emoji" alt=":sweat_smile:"> My apologies, I meant to include a description of the error:</p>
<pre><code class="lang-auto">line 101, in &lt;lambda&gt; h.register_hook(lambda x: x.clamp(min=-10, max=10))
AttributeError: 'NoneType' object has no attribute 'clamp'
</code></pre>
<p>It seems the tensor being passed to the lambda is <code>None</code></p>

username: albanD
cooked: <p>How do you use <code>h</code> later on?<br>
<code>None</code> means that it is all 0s in this context. But I am surprised the hook actually gets called with None…</p>

username: buchholzmd
cooked: <p>That’s good to know. So an uninitialized tensor defaults to <code>None</code>?</p>
<p>I seem to have a work around by using:</p>
<pre><code class="lang-auto">for h in hidden:
    if h.requires_grad:
        h.register_hook(lambda x: x.clamp(min=-10, max=10) if x is not None else x)
</code></pre>
<p>Upon examining the contents of the tensors, they contain non-zero values.</p>