Using `autograd.functional.jacobian`/`hessian` with respect to `nn.Module` parameters

Josh: <p>I was pretty happy to see that computation of Jacobian and Hessian matrices are now built into the new <code>torch.autograd.functional</code> API which avoids laboriously writing code using nested for loops and multiple calls to <code>autograd.grad</code>. However, I have been having a hard time understanding how to use them when the independent variables are parameters of an <code>nn.Module</code>. For example, I would like to be able to use <code>hessian</code> to compute the Hessian of a loss function w.r.t. the model’s parameters. If I don’t use an <code>nn.Module</code> I can successfully compute the Hessian using</p>
<pre><code class="lang-python">import torch

x = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
y = torch.tensor([1.0, 2.0])

def compute_z(a, b):
    output = (x * a.unsqueeze(0) ** 3 + x * b.unsqueeze(0) ** 3).sum(dim=1)
    z = ((output - y)**2).mean()
    return z

a = torch.tensor([1.0, -1.0, 2.0], requires_grad=True)
b = torch.tensor([2.0, -2.0, -1.0], requires_grad=True)
params = (a, b)

hessians = torch.autograd.functional.hessian(compute_z, params, strict=True)

param_names = ('a', 'b')
for d_name, d_hessians in zip(param_names, hessians):
    for dd_name, dd_hessian in zip(param_names, d_hessians):
        print(f'dz/d{dd_name}d{d_name} = \n{dd_hessian}\n')
</code></pre>
<p>which works as I would expect</p>
<pre><code class="lang-nohighlight">dz/dada = 
tensor([[ 963.,  198.,  972.],
        [ 198., -801., 1296.],
        [ 972., 1296., 9108.]])

dz/dbda = 
tensor([[ 612.,  792.,  243.],
        [ 792., 1044.,  324.],
        [3888., 5184., 1620.]])

dz/dadb = 
tensor([[ 612.,  792., 3888.],
        [ 792., 1044., 5184.],
        [ 243.,  324., 1620.]])

dz/dbdb = 
tensor([[4068., 3168.,  972.],
        [3168., 2052., 1296.],
        [ 972., 1296., -909.]])

</code></pre>
<p>But when I encapsulate this within an <code>nn.Module</code> I can’t get it to work. Among other things I’ve tried</p>
<pre><code class="lang-python">import torch

class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.a = torch.nn.Parameter(torch.tensor([1.0, -1.0, 2.0]))
        self.b = torch.nn.Parameter(torch.tensor([2.0, -2.0, -1.0]))

    def forward(self, x):
        output = (x * self.a.unsqueeze(0) ** 3 + x * self.b.unsqueeze(0) ** 3).sum(dim=1)
        return output

x = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
y = torch.tensor([1.0, 2.0])
net = Net()

def compute_z(*net_parameters):
    # Is there a proper way to set the parameters that works with hessian?
    for p_src, p_dst in zip(net_parameters, net.parameters()):
        p_dst.data = p_src.data

    output = net(x)
    z = ((output - y)**2).mean()
    return z

# strict=True raises exception, allowing non-strict for demonstration
hessians = torch.autograd.functional.hessian(compute_z, tuple(net.parameters()))

param_names = [n for n, _ in net.named_parameters()]
for d_name, d_hessians in zip(param_names, hessians):
    for dd_name, dd_hessian in zip(param_names, d_hessians):
        print(f'dz/d{dd_name}d{d_name} = \n{dd_hessian}\n')
</code></pre>
<p>but this results in all zeros</p>
<pre><code class="lang-nohighlight">dz/dada = 
tensor([[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.]])

dz/dbda = 
tensor([[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.]])

dz/dadb = 
tensor([[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.]])

dz/dbdb = 
tensor([[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.]])
</code></pre>
<p>Could someone explain what I’m doing wrong here and if there is some way to utilize <code>autograd.functional.hessian</code> to compute the hessian w.r.t. a module’s parameters?</p>

Josh: <p>So I have a better understanding why this is failing. I think that in order to avoid side-effects when calling the <code>autograd.functional.jacobian</code> function (which is used within <code>hessian</code>), the devs decide to perform a no-op on each of the input tensors and send resulting tensor to the function. The problem with this is that all <code>nn.Parameters</code> are necessarily leaf nodes, so any attempt to assign or copy data into the network parameters within the function necessarily breaks the edge in the computation graph pointing back to the inputs. Furthermore, since the inputs sent in aren’t the same objects that were passed to <code>.jacobian</code> then you end up never including the inputs to <code>compute_z</code> in the computation graph.</p>
<p>The only semi-solution I’ve found is to simply remove/comment out <a href="https://github.com/pytorch/pytorch/blob/55e225a2dc1529d9c68d5f8b333b155bd5b5b334/torch/autograd/functional.py#L423" rel="noopener nofollow ugc">this line</a> from the pytorch implementation.  This just removes the no-op but unfortunately that’s obviously not viable <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=9" title=":frowning:" class="emoji" alt=":frowning:"> My guess is this may also have other consequences, for example if there are multiple branching and merging paths in the computation graph then we may get the correct results?</p>

Ren_Pang: <p>I’m also working on this issue when calculating the influence function.<br>
My current solution is to follow this workaround:</p><aside class="quote quote-modified" data-post="1" data-topic="70076">
  <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/pinocchio/48/22093_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-does-one-have-the-parameters-of-a-model-not-be-leafs/70076">How does one have the parameters of a model NOT BE LEAFS?</a> 
  </div>
  <blockquote>
    I want the parameters of the model to NOT be leafs. I also want the parameters to be part of the computation graph and I am dynamically changing the parameters with code. So I want: 

Parameters to NOT be leafs (but still be held inside a nn.Module so that I can use the dynamically changing parameters as part of the computation graph).
I want the “updating model” to be part of the computation graph at any time but it seems setting the parameters to leaf is stopping the flow of gradients

How can…
  </blockquote>
</aside>
<p>
So I write this sample code (which works as a workaround):</p>
<pre data-code-wrap="python3"><code class="lang-python3">import torch
import torch.nn as nn

_input=torch.randn(32,3)
layer = nn.Linear(3,4)
criterion=nn.CrossEntropyLoss()

weight = layer.weight

def func(weight):
    del layer.weight
    layer.weight=weight
    return criterion(layer(_input), torch.zeros(len(_input),dtype=torch.long))

torch.autograd.functional.hessian(func,weight)
</code></pre>
<p>The drawback of this workaround is it’s <strong>Not Safe</strong>: You have to delete the attribute <code>weight</code> in the <code>func</code> <strong>everytime</strong>, and can’t be written just before <code>func</code> for only onece. Otherwise, the hessian will all be 0.</p>
<p>I would appreciate if there is any more elegant way. And if not, maybe we need to think about a modification for <code>torch.autograd.functional.hessian</code></p>

gngdb: <p>I wrote a general way of doing this not safe monkey patching to substitute weights into the model. The code’s <a href="https://github.com/gngdb/pytorch-minimize/blob/master/pytorch_minimize/optim.py#L106-L122" rel="noopener nofollow ugc">here</a> as part of a way to call <code>scipy.optimize.minimize</code> as a PyTorch optimizer. It’s the same except for doing a deepcopy of the Module to avoid any state problems on subsequent iterations. It’s not good for efficiency.</p>

shengchao-lin: <p>I have encountered the same issue when computing hessian or <code>vhp</code> with respect to network parameter. Hope there’s an elegant way like <code>.grad</code>.</p>

shengchao-lin: <p>This heavily relies on the fact that you know the attributes <code>layer.weight</code>. This is not possible when you only have general access like <code>layer.parameters()</code>.</p>

Ren_Pang: <p><a class="mention" href="/u/gngdb">@gngdb</a><br>
<a class="mention" href="/u/shengchao-lin">@shengchao-lin</a></p>
<p>See <a href="https://github.com/pytorch/pytorch/issues/49171#issuecomment-934055267" class="inline-onebox" rel="noopener nofollow ugc">NN Module functional API · Issue #49171 · pytorch/pytorch · GitHub</a><br>
for a better solution, which shall be released with pytorch 1.11.</p>