How to get gradients of each node in the network (not weights)

SeoHyeong: <p>.grad() method returns gradient values for each weight in the network. However, how can I get gradient values for each node in the network? Or is it safe to simply add gradients for each weight that corresponds to a specific node?</p>

ptrblck: <p>I’m not sure I understand the question properly.<br>
Could you explain your definition of <code>weight</code> vs. <code>node</code>, please? <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

SeoHyeong: <p>Yes, of course! Sorry for not being clear.</p>
<p>Let’s say I have a super simple network that has an in_feature dimension of 3 and an out_feature dimension of 5 and in Pytorch, this calculation would be done with a matrix multiplication (<code>nn.Linear(3, 5)</code>)</p>
<pre><code class="lang-auto">model = nn.Sequential(
    nn.Linear(3, 5)
)

loss.backward()
</code></pre>
<p>Then, calling .<code>grad()</code> on weights of the model will return a tensor sized 5x3 and each gradient value is matched to each weight in the model. Here, I mean weights by connecting lines in the figure below.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/4/646a9d2c85d34157c4b6c51adedc9a10cb9b2083.png" data-download-href="https://discuss.pytorch.org/uploads/default/646a9d2c85d34157c4b6c51adedc9a10cb9b2083" title="Screen Shot 2022-03-10 at 6.47.17 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/4/646a9d2c85d34157c4b6c51adedc9a10cb9b2083_2_593x499.png" alt="Screen Shot 2022-03-10 at 6.47.17 PM" data-base62-sha1="ekkajEgKZxAvmqC0bx92Ou2laBJ" width="593" height="499" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/4/646a9d2c85d34157c4b6c51adedc9a10cb9b2083_2_593x499.png, https://discuss.pytorch.org/uploads/default/optimized/3X/6/4/646a9d2c85d34157c4b6c51adedc9a10cb9b2083_2_889x748.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/4/646a9d2c85d34157c4b6c51adedc9a10cb9b2083.png 2x" data-dominant-color="EEEEEE"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2022-03-10 at 6.47.17 PM</span><span class="informations">1158×976 89.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>However, I was wondering how to get gradients for each of the 3 in_feature dimensions, which are the yellow dots (nodes). Would summing the 5x3 sized gradient matrix into the dimension of 3 be a safe way to get them?</p>

AlphaBetaGamma96: <p>I was wondering if you might be able to get this via using a <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=full%20backward#torch.nn.Module.register_full_backward_hook" rel="noopener nofollow ugc">backward hook</a>? Because it seems like you’re looking for the <code>grad_input</code> of a layer? Does that work <a class="mention" href="/u/ptrblck">@ptrblck</a> ?</p>

SeoHyeong: <p>Just to clarify that my question wasn’t specifically toward on gradients w.r.t the input but it was for those neurons in the network in general. So if the network was deeper than the example above, I’d like to get gradients of each and every neuron of the network! Any help would be appreciated <img src="https://discuss.pytorch.org/images/emoji/apple/slightly_smiling_face.png?v=12" title=":slightly_smiling_face:" class="emoji" alt=":slightly_smiling_face:" loading="lazy" width="20" height="20"></p>

AlphaBetaGamma96: <p>yes, but surely the neurons are just inputs of a given layer? So, if all your layers are just <code>nn.Module</code> objects you should be able to get the gradient with respect to each neuron via the backward hook?</p>
<p>Note: the backward hook returns the gradients for each input in your batch too, so you may need to take the mean over the batch dim to get the shape you want.</p>

ptrblck: <p>I agree with <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a>’s suggestions and think backward hooks should work to get the gradient for the input activations.</p>

AlphaBetaGamma96: <p>Make sure to use <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=full_backward#torch.nn.Module.register_full_backward_hook" rel="noopener nofollow ugc">register_full_backward_hook</a> and not <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_backward_hook" rel="noopener nofollow ugc">register_backward_hook</a> as <code>register_backward_hook</code> is deprecated and gives unexpected behaviour!</p>

SeoHyeong: <p>I was able to get gradient information w.r.t every layer in my model with <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=full_backward#torch.nn.Module.register_full_backward_hook" rel="noopener nofollow ugc">register_full_backward_hook </a>. Thank you <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>I have an additional question on the behavior of <code>register_full_backward_hook</code>.</p>
<p>Below is the printed output of my code with <code>register_full_backward_hook</code> and I expect input grad shape and output grad shape for <code>nn.Linear(512, 512)</code> would be <code>torch.Size([12, 512])</code> and <code>torch.Size([12, 512])</code> each (12 is the size of the mini-batch). However, it’s giving me another dimension of 30.</p>
<pre><code class="lang-auto">hook triggered on Linear(in_features=512, out_features=512, bias=False)
input grad shape: torch.Size([12, 30, 512])
output grad shape: torch.Size([12, 30, 512])
</code></pre>
<p>Do you have any idea where that 30 comes from?</p>

ptrblck: <p>Are you passing an input of <code>[12, 30, 512]</code> to this layer as seen here?</p>
<pre><code class="lang-python">def hook(module, grad_input, grad_output):
    print(grad_input[0].shape)
    print(grad_output[0].shape)
    
lin = nn.Linear(512, 512, bias=False)
lin.register_full_backward_hook(hook)

x = torch.randn(12, 30, 512, requires_grad=True)
out = lin(x)
out.mean().backward()
</code></pre>
<p>Output:</p>
<pre><code class="lang-python">torch.Size([12, 30, 512])
torch.Size([12, 30, 512])
</code></pre>

SeoHyeong: <p>Oh yes! Just checked the input dimension size and turns out the layer was getting an input of <code>torch.Size([12, 30, 512])</code>. Thank you so much!</p>