Backward hook uncaught in grad of grad

username: jusjusjus
cooked: <p>Hi pytorch friends!</p>
<p>I’m trying to implement fast gradient penalty using forward and backward hooks but found that for gradients of gradients, hooks show a slightly aberrant behavior.  Hopefully, you can help me find where I go wrong.</p>
<p>I tried to construct a minimal example that shows the behavior.  Let me start by hooking in a simple linear model with one parameter.</p>
<pre><code class="lang-python">import numpy as np
import torch
from torch import nn
from torch import autograd

torch.manual_seed(42)

class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(1, 1, bias=False)
        self.lin.weight = nn.Parameter(torch.tensor([[2.0]]))

    def forward(self, x):
        # Switch comments in the following lines for an aberrant behavior.
        return self.lin(x)
        # return torch.log(torch.exp(self.lin(x)))

model = Model()

# Add hooks to construct per-example gradients
gradients = {}
activations = {}
def forward_hook(layer, inputs, outputs):
    activations[layer] = inputs[0].detach()

def backward_hook(layer, grad_input, grad_output):
    A = activations[layer]
    B = grad_output[0].detach()
    gradients[layer] = torch.einsum("ni,nj-&gt;nij", B, A)

model.lin.register_forward_hook(forward_hook)
model.lin.register_backward_hook(backward_hook)
</code></pre>
<p>The backward hook computes the gradient specifically for a layer of type <code>nn.Linear</code>.  Let’s feed the model with a batch of 2 input examples, and compute parameter gradients.  We also check that the per-example gradients are computed correctly.  But we’ll comment that out later.</p>
<pre><code class="lang-python">x = torch.FloatTensor([2, 1])[..., None]
x.requires_grad_(True)
loss = model(x).sum()
# We'll comment this out later
loss.backward()
for layer, g in gradients.items():
    assert torch.allclose(layer.weight.grad, g.sum(dim=0))
# .. until here.
</code></pre>
<p>So the hooks seem to work, and I was pretty happy with myself having successfully stolen the trick of <a class="mention" href="/u/yaroslav_bulatov">@Yaroslav_Bulatov</a> and his colleagues.</p>
<p>However, I need to actually compute gradient of <code>loss</code> with respect to the input data <code>x</code>, and construct a cost function from that.</p>
<pre><code class="lang-python">grads = autograd.grad(
    outputs=loss, inputs=x,
    retain_graph=True, create_graph=True, only_inputs=True
)[0]
grad_loss = (grads ** 2).sum()
model.zero_grad()
grad_loss.backward()
</code></pre>
<p>Unfortunately the test fails now:</p>
<pre><code class="lang-python">grad_loss = (grads ** 2).sum()
model.zero_grad()
grad_loss.backward()
for layer, g in gradients.items():
    assert layer.weight.grad == 8, layer.weight.grad
    assert g.sum(dim=0) == 3, g
</code></pre>
<p>The situation gets even more confusing when you change the comments in <code>Model.forward</code>.  Note that we take log(exp(x)) which should do anything, but subsequently, the hooked gradients are all zero.</p>
<p><a href="https://colab.research.google.com/drive/15ocWwMld2-Q5VJb6NKqJxwtXeFwKczDF" rel="nofollow noopener">Here’s</a> a notebook.</p>
<p>Please help <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"></p>

username: albanD
cooked: <p>Hi,</p>
<p>As mentionned in the doc for <code>.register_backward_hook()</code>, they are fairly broken right now and will give aberrant behavior :’( We are working on fixing it but it proves quite challenging.</p>
<p>In particular, the <code>grad_input</code> value will be wrong as soon as your <code>nn.Module </code> does more than one operation.</p>
<p>If you have a single Tensor for input and outputs, you should use <code>tensor.register_hook()</code> that is called on a Tensor in combination with the <code>nn.Module.register_forward_hook()</code> to get the behavior you want:</p>
<p>This would look like this (I did not run the code so there may be typos)</p>
<pre><code class="lang-auto"># Add hooks to construct per-example gradients
gradients = {}
activations = {}
def forward_hook(layer, inputs, outputs):
    activations[layer] = inputs[0].detach()

    output_grads = []
    def hook_output(grad_output):
        output_grads[0] = grad_output.detach()

    def hook_input(grad_input):
        A = activations[layer]
        B = output_grads[0]
        gradients[layer] = torch.einsum("ni,nj-&gt;nij", B, A)

    inputs[0].register_hook(hook_input)
    outputs[0].register_hook(hook_output)

model.lin.register_forward_hook(forward_hook)
</code></pre>

username: jusjusjus
cooked: <p>I explored your solution, but without much luck so far.  Indeed, there’s also a typo: <code>outputs[0].register...</code> -&gt; <code>outputs.register...</code>.   Thanks though; let me how I can help with a module fix.</p>

username: saurabheights
cooked: <aside class="quote no-group" data-username="albanD" data-post="2" data-topic="72026">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p><code>activations[layer] = inputs[0].detach()</code></p>
</blockquote>
</aside>
<p>Shouldn’t activations of layer be the <code>outputs</code> variable?<br>

P.S. Since we are using the word activation, it should require applying relu or other activation.
(Sorry, ignore this).</p>