Track Intermediate Gradient within Module

ludwigwinkler: <p>Hi everybody,</p>
<p>I want to track intermediate gradients in the computational graph.</p>
<p>There have been related questions on this as in</p><aside class="quote" data-post="1" data-topic="19565">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/blade6570/48/3508_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/register-hook-and-find-the-maximum-gradient-of-non-leaf-variable/19565">Register_hook and find the maximum gradient of non leaf variable</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    I have put the register_hook() for a non_leaf variable and printing the gradient. How to access each element in the gradient tensor. 
The syntax I am using is as follows: 
y.register_hook(print)
  </blockquote>
</aside>
<aside class="quote quote-modified" data-post="1" data-topic="94">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/kalamaya/48/381_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/why-cant-i-see-grad-of-an-intermediate-variable/94">Why cant I see .grad of an intermediate variable?</a> 
  </div>
  <blockquote>
    Hi! I am loving this framework… 
Since Im a noob, I am probably not getting something, but I am wondering why I cant get the gradient of an intermediate variable with .grad? Here is an example of what I mean: 
xx = Variable(torch.randn(1,1), requires_grad = True)
yy = 3*xx
zz = yy**2
zz.backward()
xx.grad # This is ok
yy.grad # This gives 0! 
zz.grad # This should give 1! 

So I get the correct result for xx.grad, but why does yy.grad show 0, as does zz.grad? How can I get the yy.grad value in t…
  </blockquote>
</aside>

<p>Yet the solution to both problems were applied to fairly simple and straight forward computation graphs.</p>
<p>I’m interested in wrapping the tracking of the intermediate gradients in an optimizer class such that I collect the intermediate non-leaf gradients and process them with more complicated preconditioning methods in the fancy_optim.step() function.</p>
<p>Down below is a minimum working example in which I want to store the gradients of the class variable self.intermediate gradient.</p>
<p>The tracking of parameters is relatively straight forward and I’m aware that the two matrix multiplies can be combined to one.</p>
<p>Is there some way to save the intermediate gradients in an optimizer class through hooks?<br>
Or maybe some other way to track intermediate gradients through the optimizer class?</p>
<p>PS: Is my understanding correct that the self.state variable in the optimizer class stores all relevant values for the optimization of the parameters?</p>
<pre><code class="lang-auto">import torch
import torch.nn.functional as F

class ExampleLayer(torch.nn.Module):

	def __init__(self, _dim_input, _dim_intermediate, _dim_output):

		super().__init__()

		self.matrix1 = torch.nn.Parameter(torch.randn(_dim_input, _dim_intermediate))

		# I want to store and track these gradients
		self.intermediate_gradient = torch.randn(_dim_input, _dim_output).requires_grad_()

		# But only this and matrix1 are parameters
		self.matrix2 = torch.nn.Parameter(torch.randn(_dim_intermediate, _dim_output))

	def forward(self, _input: torch.Tensor):

		out = torch.mm(_input, self.matrix1)

		self.intermediate_gradient = torch.mm(out, self.matrix2)
		self.intermediate_gradient.retain_grad()

		return self.intermediate_gradient

dim_input = 21
dim_intermediate = 11
dim_output = 31

x = torch.randn(100, dim_input) # Batch size of 100

layer = ExampleLayer(dim_input, dim_intermediate, dim_output)
optim = torch.optim.Adam(layer.parameters())

saved_grads = []

def save_grad():

	def savegrad_hook(grad):
		print('inside savegrad_hook')
		saved_grads.append(grad)
	print('inside function save_grad')
	return savegrad_hook

layer.intermediate_gradient.register_hook(save_grad())

for epoch in range(3):
	out = layer(x)
	F.mse_loss(out, torch.ones_like(out)).backward()

	print(layer.intermediate_gradient.grad)

print(saved_grads)

# print(layer.intermediate_gradient.grad)
</code></pre>
<p>Thanks!</p>

ludwigwinkler: <p>Okay, I was able to solve the problem by realizing that the register_hook command for non-leaf nodes has to be placed precisely between forward and backward pass.</p>
<p>In comparison the register_forward_hook and register_backward_hook can be initialized before the forward pass.</p>
<p>While it makes a lot of sense in hindsight, maybe it could be noted in the docs?</p>