Recording gradients in RNNs at each point in time

mllera: <p>Hi,</p>
<p>I am trying to record the values of the gradients as they propagate in time through an RNN. Initially I thought this would be easily accomplished using  the <code>register_hook</code> function by calling it to each parameter, yet I now realize that for some reason, the hook is not called at each step in time, contrary to what I had understood. Take for example the following code:</p>
<pre><code class="lang-python">import torch

x = torch.randn((1, 1))
w = torch.ones((2, 1), requires_grad=1)

z = w * x
z *= w
z *= w
z *= w

def print_grad(grad):
    print(grad)

h = z.register_hook(print_grad)
z.sum().backward()

</code></pre>
<p>This produces as output:</p>
<pre><code class="lang-bash">tensor([[1.],  [1.]])
</code></pre>
<p>Instead of the same, but repeated 4 times.</p>
<p>What am I missing here? If nothing, is there a way to record the gradients at each point in time?</p>
<p>Thanks in advance!</p>

albanD: <p>The hook will give you the gradient at the point where the hook is registered.<br>
If you want to get it after every computation, you need to register one after every computation.</p>

mllera: <p>Oh, I see. I guess for feed-forward models that’s good enough. But for recurrent networks it makes tracing gradients backward in time a bit hard.</p>
<p>Thanks for the answer!</p>

albanD: <p>Well if you have a single call in a for-loop, you will still have one call to the hook per time it was applied. So that works for rnn as well <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<pre><code class="lang-auto">def print_grad(grad):
    print(grad)

for i in range(4):
  z *= w
  z.register_hook(print_grad)
</code></pre>
<p>That will register one hook per iteration of the loop.</p>

mllera: <p>That would certainly work, but I still have to modify the code in all the RNN models to save the gradients, which is not what I want. The recording and the model itself should be separate. It’s also not what the documentation implies. Taken directly from there:</p>
<pre><code class="lang-python">def register_hook(self, hook):
        r"""Registers a backward hook.

        The hook will be called every time a gradient with respect to the
        Tensor is computed. The hook should have the following signature::

            hook(grad) -&gt; Tensor or None
...
</code></pre>
<p>Which implies that it should be called at each time step, since we compute a gradient w.r.t. the parameters for each of them. Yet when I do this, I only get one gradient.</p>
<p>Maybe someone want’s the gradient for the whole sequence, but that can be obtained using the <code>register_backward_hook</code> method in the Module class.</p>

albanD: <p>Every time you do <code>z = foo(z)</code> in your RNN, you most likely create a new Tensor that is then assigned with the python variable named <code>z</code>. But the original Tensor contained in <code>z</code> and the new one have nothing in common. So if you add a hook to the second one, you won’t get gradient for the first one.</p>
<p>This comment refers to the case where you call <code>.bakward()</code> multiple times.</p>

mllera: <p>Yes, but I am trying to compute the gradient w.r.t. the parameters, which are tensors that do not change and are repeated across time steps. Also I only call backwards one time. From what I read in the docs, this should give me a gradient for each time step and parameters (i.e. input and recurrent weights and biases), unless I have a bug or misunderstood something…</p>

albanD: <p>If you use a Tensor multiple times, the gradient associated to it is the sum of the gradients corresponding to each use.<br>
If you want to get the gradient for a single use, you need to make a temporary tensor corresponding to that single use.</p>

metro.smiles: <p>So interestingly enough, the gradient returned in every time step by creating this hook is the same. I’m guessing it’s the sum total of all gradients of z over every time step. Can someone confirm this?</p>

albanD: <p>If the Tensor is the same, then yes.</p>

metro.smiles: <p>And just of curiosity, there’s no way to extract the gradients for every time step separately (rather than deal with the sum) ?</p>

albanD: <p>As mentioned above, you can do that by having a different Tensor (with same content) for each step.</p>