Calculating the Hessian of loss function wrt torch network parameters

semihcanturk: <p>Is there an efficient way to compute second order gradients (at least a partial Hessian) of a loss function with respect to the parameters of a network using PyTorch autograd? How <code>torch.autograd.functional.hessian(func, inputs, ...)</code> works doesn’t play nice at all with Torch modules after all, since a standard loss function does not take the network parameters themselves as inputs, and operates on the network itself.</p>
<p>Two leads I have on this are the following, but neither really solve the problem:</p>
<ol>
<li>
<p>A similar problem is discussed here in the forums: <a href="https://discuss.pytorch.org/t/using-autograd-functional-jacobian-hessian-with-respect-to-nn-module-parameters/103994" class="inline-onebox">Using `autograd.functional.jacobian`/`hessian` with respect to `nn.Module` parameters</a>. It’s a bit outdated now perhaps, so I don’t know if there are more viable solutions at the moment.</p>
</li>
<li>
<p>I found a quite cumbersome possible workaround in <a href="https://stackoverflow.com/questions/64024312/how-to-compute-hessian-matrix-for-all-parameters-in-a-network-in-pytorch" class="inline-onebox" rel="noopener nofollow ugc">machine learning - How to compute hessian matrix for all parameters in a network in pytorch? - Stack Overflow</a>, which suggests writing a wrapper function that:<br>
a) takes in flattenned all network parameters,<br>
b) unflattens them inside,<br>
c) and then basically mimicks a forward pass and computes the loss,<br>
all in order to make it play nice with <code>torch.autograd.functional.hessian()</code>. It <em>could</em> work, but I feel there has to be a better way…</p>
</li>
</ol>
<p>Does anyone have a better way to work around this? Thank you!</p>

Mari: <p>Hi,</p>
<p>I am trying to compute Hessian of a loss function w.r.t  the parameters of a network using PyTorch autograd. Can you please share your code here or provide me an example how I can do it?</p>
<p>Thanks,<br>
Maryam</p>

y12uc231: <p>On the same boat, haven’t gotten a great way to do it yet. Please share if anyone has. Thanks</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/y12uc231">@y12uc231</a>, <a class="mention" href="/u/mari">@Mari</a> and <a class="mention" href="/u/semihcanturk">@semihcanturk</a>,</p>
<p>If you want to compute the Hessian of a network, you might want to have a look at the functorch library (which now comes packaged with the latest install of PyTorch). A brief example I found on their github page is here: <a href="https://github.com/pytorch/functorch/issues/989" rel="noopener nofollow ugc">https://github.com/pytorch/functorch/issues/989</a>.</p>
<p>In general, you can do something like this,</p>
<pre><code class="lang-auto">from functorch import make_functional, jacfwd, jacrev, vmap

net = Model() #model instance
fnet, params = make_functional(net) #functorch requires a functional form of the model (fnet) with parameters (params) as an input to the model as well.

per_sample_hessian = vmap(jacfwd(jacrev(fnet, argnums=0), argnums=0), in_dims=(None, 0))(params, x)
</code></pre>
<p>where <code>x</code> is our input to the model and we’re calculating the hessian with respect to our network’s parameters via forward-over-reverse Hessian calculation method for efficiency. The example also uses functorch’s vmap to vectorize our calculation over the batch dimension, which is critical in calculating higher-order calculations efficiently!</p>
<p>It’s quite a lot to take in so I recommend you read their notebook on Hessian calculations (among the rest of the functorch documentation) here: <a href="https://pytorch.org/functorch/stable/notebooks/jacobians_hessians.html" class="inline-onebox" rel="noopener nofollow ugc">Jacobians, Hessians, hvp, vhp, and more: composing functorch transforms — functorch 1.13 documentation</a></p>