Calculate second derivative related to preactivations

username: Keru
cooked: <p>Let s denote preactivation units (before ReLU, after Linear), then I can get gradients of preactivations dE/ds through register_backward_hook. But how I can get gradients of f(dE/ds) w.r.t. weights, where f is another function (e.g. norm)? It seems that grad_out in the hook does not require grad and does not have a grad_fn.</p>
<p>This is similar to calculating second derivative, except that I’m considering f(dE/ds) instead of f(dE/dW). I know this is somehow abnormal, but I’m trying some new ways related to Hessian matrix. Just wondering whether this is possible…</p>
<p>Thanks</p>

username: albanD
cooked: <p>Hi,</p>
<p>First i would advise you not to use hooks on nn.Modules but add them to Tensors: <code>s.register_hook(your_fn)</code>. In particular, as you can see in the doc, register_backward_hook is discouraged.</p>
<p>Now in your hook. The gradient will require gradients only if you backardprop with <code>create_graph=True</code>.</p>

username: Keru
cooked: <p>Thanks a lot for your reply!</p>
<p>I tried <code>s.register_hook(fn)</code> and <code>create_graph = True</code>. Problem is that both input and output of hook function <code>fn</code> still have <code>requires_grad = False</code> and no <code>grad_fn</code>. So I can’t calculate gradient again… (If I register_hook for weights, they do have <code>requires_grad = True</code> and <code>grad_fn</code>)</p>
<p>I know that we can use <code>torch.autograd.grad</code> to calculate higher order gradients for <code>w</code>, but preactivation <code>s</code> is non-leaf variable. Then <code>grad(loss, s)</code> doesn’t have <code>requires_grad</code> nor <code>grad_fn</code>.</p>
<p>It seems that we can only calculate higher order gradients for leaf variables?</p>

username: albanD
cooked: <p>For intermediary results, they need to require gradients for you to be able to compute gradients on this. If they don’t during the forward pass, you need to add it with <code>.requires_grad_()</code> before continuing the forward pass.</p>

username: Keru
cooked: <p>In the following toy example, I register hooks for each variable, but only hook grad for <code>x</code> and <code>w</code> print <code>True</code> and <code>s</code> will print <code>False</code>.</p>
<pre><code class="lang-auto">def hk(x):
    print(x.requires_grad)

x = torch.randn(16, 5, requires_grad=True)
x.register_hook(hk)
w = torch.randn(5, 3, requires_grad=True)
w.register_hook(hk)

s = torch.mm(x,w)
s.requires_grad_()
s.retain_grad()
s.register_hook(hk)

y = s.sum()
y.backward(create_graph=True)
</code></pre>
<p>I further set <code>s.retain_grad()</code>, but only <code>x.grad</code> and <code>w.grad</code> have <code>grad_fn</code>, and <code>s.grad</code> doesn’t.</p>
<p>It seems that I can only get grad values of intermediate variables rather than further processing them.</p>

username: albanD
cooked: <p>Let me add some comments to your code:</p>
<pre><code class="lang-python">import torch

def get_hook(msg):
    def hk(x):
        print("Does {}'s grad requires grad?".format(msg))
        print(x.requires_grad)
    return hk

x = torch.randn(16, 5, requires_grad=True)
x.register_hook(get_hook("x"))
w = torch.randn(5, 3, requires_grad=True)
w.register_hook(get_hook("w"))

s = torch.mm(x,w)
s.register_hook(get_hook("s"))

y = s.sum()
# use retain_graph=True to be able to call backward multiple times
y.backward(retain_graph=True)
print("All 3 hooks are called and all print False")

# Asking for graph to be created and set the grad output to requires_grad=True
go = torch.ones((), requires_grad=True)
y.backward(go, create_graph=True, retain_graph=True)
print("All 3 hooks are called and all print True")

# Asking for graph to be created but the grad output has requires_grad=False
y.backward(create_graph=True, retain_graph=True)
print("All 3 hooks are called and s print False but x,w print True")
# This is because go does not require gradients, so the gradient of s doesn't either
# But in the mm, x is used to compute the gradient of w, so the gradient of w requires gradients.

</code></pre>