Implementing calculation of the Laplacian

username: AlphaBetaGamma96
cooked: <p>Hi All,</p>
<p>I have a quick question regarding how to implement the Laplacian of the output of a network efficiently. I’ve managed to implement a method for calculating it, however, I’m pretty sure the way I’m doing it is a pretty inefficient way.</p>
<p>What I mean by the Laplacian of the output of the network is, let’s say I have a simple feed-forward network, <code>y = model(x)</code> and I wish to calculate <code>Sum_{i} d2y_dx_{i}2</code> for all inputs samples.</p>
<p>My current code for calculating the Laplacian for an N-dimensional input <code>x</code> is</p>
<pre><code class="lang-auto">y = model(x) #where model is an R^N to R^1 function

laplacian = torch.zeros(x.shape[0]) #array to store values of laplacian

for i, xi in enumerate(x):
    hess = torch.autograd.functional.hessian(model, xi.unsqueeze(0), create_graph=True)
    laplacian[i] = torch.diagonal(hess.view(N, N) offset=0).sum()
</code></pre>
<p>I’ve noticed for large input samples (which I need) the memory usage increases drastically. I would assume this is because I create a graph for each call of <code>torch.autograd.functional.hessian</code>. I need this because I need to differentiate the Laplacian with respect to the network parameters in my loss function.</p>
<p>My question is, would it be possible to perhaps reduce the memory usage on calculating the hessian or perhaps re-use the same graph but just for different inputs? (if that’s the correct way of phrasing it?)</p>
<p>Many thanks in advance!</p>

username: albanD
cooked: <aside class="quote no-group" data-username="AlphaBetaGamma96" data-post="1" data-topic="99644">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3da27b/48.png" class="avatar"> AlphaBetaGamma96:</div>
<blockquote>
<p>re-use the same graph but just for different inputs?</p>
</blockquote>
</aside>
<p>The whole memory usage form the graph comes from intermediary results. So re-using the graph wouldn’t help.</p>
<p>I’m afraid that extracting the diagonal of the Hessian cannot be done more efficiently in general.<br>
You could reduce the memory usage by doing a custom function that drops the non-diagonal elements earlier but that won’t be a huge change.</p>
<p>You might want to check packages like this one <a href="https://backpack.pt/">https://backpack.pt/</a> that provides approximate method to get the diagonal elements. If you are happy with an approximate value for your use case, it will be the fastest you can do!</p>

username: AlphaBetaGamma96
cooked: <p>Thanks for the quick response! On implementing a custom function, I did have a look at creating custom autograd functions which have forward and backward methods. I was wondering if it were at all possible to implement a backward of backward method in an autograd function?</p>
<p>What I mean by backward of backward would be an analytical expression for the 2nd order gradient of the loss function with respect to the output of the final layer. Because the final layer of my network is a determinant function, explicitly writing out those derivatives and calling them directly would be far more efficient than calling autograd and letting it solve it itself.</p>

username: albanD
cooked: <aside class="quote no-group" data-username="AlphaBetaGamma96" data-post="3" data-topic="99644">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3da27b/48.png" class="avatar"> AlphaBetaGamma96:</div>
<blockquote>
<p>I was wondering if it were at all possible to implement a backward of backward method in an autograd function?</p>
</blockquote>
</aside>
<p>The engine expects to only compute gradients. So you do have to do two backwards right now.</p>
<p>You can write these formulas for sure but it might be tricky to be re-use the current engine for this <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=12" title=":confused:" class="emoji" alt=":confused:" loading="lazy" width="20" height="20"></p>

username: rfmiotto
cooked: <p>Hi, AlphaBetaGamma96,</p>
<p>I am having the same issue here. Memory just blows up when computing the Laplacian.<br>
Did you manage to find an efficient implementation for the Laplacian?<br>
If so, could you please share it?</p>
<p>Best</p>

username: AlphaBetaGamma96
cooked: <p>Hi <a class="mention" href="/u/rfmiotto">@rfmiotto</a>,</p>
<p>You can do this quite efficiently within the <code>torch.func</code> namespace, which requires version torch2.1+. Here’s an example below for the laplacian with respect to an input <code>x</code> below,</p>
<pre><code class="lang-auto">import torch
from torch.func import hessian, vmap, functional_call

net = Model(*args, **kwargs) #our model

x = torch.randn(batch_size, num_inputs) #inputs
params = dict(net.named_parameters()) #params (need to be an input for torch.func to work)

#functionalize version (params are now an input to the model)
def fcall(parms, x):
  return functional_call(net, params, x)

def compute_laplacian(params, x):
  _hessian_wrt_x = hessian(fcall, argnums=1)(params, x) #forward-over-reverse hessian calc.
  _laplacian_wrt_x = _hessian_wrt_x.diagonal(0,-2,-1) #use relative dims for vmap (function doesn't see the batch dim of the input)
  return _laplacian_wrt_x

#We define a laplacian func for a single function, then vectorize over the batch.
laplacian = vmap(compute_laplacian, in_dims=(None,0))(params, x) 
</code></pre>

username: rfmiotto
cooked: <p>Thanks for sharing!<br>
I’ll test your version soon to see if I can run the training without hitting the memory limit.</p>
<p>In the meantime, I’d like to share the approach I ended up taking to get around the memory problem: running using finite differences. I do a single forward pass, but triple the input size due to the stencil to calculate the numerical differentiation.<br>
One can find the implementation <a href="https://github.com/las-unicamp/siren/blob/main/src/vector_ops.py" rel="noopener nofollow ugc">here</a>.</p>