How to avoid sum from ‘autograd.grad’ output in Physics Informed Neural Network?

username: Isabella_Osei
cooked: <p>Hello,<br>
I’m working on a Physics Informed Neural Network and I need to take the derivatives of the outputs w.r.t the inputs and use them in the loss function.<br>
The issue is related to the neural network’s multiple outputs. I tried to use ‘autograd.grad’ to calculate the derivatives of the outputs, but it sums all the contributions.<br>
For example, if my output ‘u’ has shape [batch_size, n_output], the derivative ‘dudx’ has shape [batch_size, 1], instead of [batch_size, n_output].<br>
Due to the sum, I can’t use the derivatives in the loss function. I tried with a for loop to calculate each derivative but the training takes forever. Do you have any idea how to solve this problem? Thanks in advance</p>

username: AlphaBetaGamma96
cooked: <p>You could have a look at using <code>torch.func.jacrev</code> and <code>torch.func.vmap</code> to compute the entire jacobian of your network.</p>
<p>Also, please share a minimal reproducible example to help explain your problem.</p>

username: Isabella_Osei
cooked: <p>Hi,<br>
thanks, you are right I should have posted some code.<br>
below the code as per my loss function contains high-order derivatives of the outputs with respect to the inputs x and y</p>
<pre><code class="lang-auto">def gradient(y, x, grad_outputs=None):
    if grad_outputs is None:
        grad_outputs = torch.ones_like(y)
    grad = torch.autograd.grad(y, [x], grad_outputs=grad_outputs, create_graph=True)[0]
    return grad


def compute_derivatives(x, y, u):
    dudx = gradient(u, x)
    dudy = gradient(u, y)

    dudxx = gradient(dudx, x)
    dudyy = gradient(dudy, y)

    dudxxx = gradient(dudxx, x)
    dudxxy = gradient(dudxx, y)
    dudyyy = gradient(dudy, y)

    dudxxxx = gradient(dudxxx, x)
    dudxxyy = gradient(dudxxy, y)
    dudyyyy = gradient(dudyyy, y)

    return dudxx, dudyy, dudxxxx, dudyyyy, dudxxyy

</code></pre>
<p>Also, I already tried vmap with a simplified version of code but it gives back the following error:</p>
<pre><code class="lang-auto">import torch
from functorch import vmap

x = torch.tensor([1.0, 2.0], requires_grad=True)
out = torch.stack([x * 2, x * 3], dim=0)

print('x:', x)
print('out:', out)


def single_gradient(out_row, x):
    grad_outputs = torch.ones_like(out_row)
    return torch.autograd.grad(out_row, [x], grad_outputs=grad_outputs, create_graph=True, retain_graph=True)[0]


batched_grad = vmap(single_gradient, (0, None))(out, x)

print('Batched Grads:', batched_grad)
</code></pre>
<p>“RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn”<br>
What do you think about this?</p>
<p>The Jacobian requires as input a function, but my input is the output with shape [batch_size,N] from the neural network.<br>
Given I do not have a function, can I still use Jacobian?</p>

username: AlphaBetaGamma96
cooked: <aside class="quote no-group" data-username="Isabella_Osei" data-post="3" data-topic="208007">
<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/i/d2c977/48.png" class="avatar"> Isabella_Osei:</div>
<blockquote>
<p>“RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn”</p>
</blockquote>
</aside>
<p>Using the <code>torch.func</code> namespace requires you to re-write how you compute the gradients entirely, you can’t mix and match between <code>torch.autograd</code> and <code>torch.func</code>, especially when using <code>torch.func.vmap</code> (at least to my knowledge)</p>
<p>I have some previous examples on the forums of how to compute gradients with a ‘functional’ approach here:</p>
<aside class="quote quote-modified" data-post="8" data-topic="156222">
  <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">
    <a href="https://discuss.pytorch.org/t/efficient-computation-of-hessian-with-respect-to-network-weights-using-autograd-grad-and-symmetry-of-hessian-matrix/156222/8">Efficient computation of Hessian with respect to network weights using autograd.grad and symmetry of Hessian matrix</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" style="--category-badge-color: #AB9364; --category-badge-text-color: #FFFFFF;" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    The documentation of the hessian function for functorch can be found 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 nightly documentation</a> 
I would advise using torch.func instead (if you’re using a more recent version of PyTorch), but something similar is possible in functorch. In torch.func an example of computing the Hessian with respect to parameters would be something like, 
import torch
from torch import nn
from torch.func import functional_call, v…
  </blockquote>
</aside>

<aside class="quote quote-modified" data-post="6" data-topic="206553">
  <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">
    <a href="https://discuss.pytorch.org/t/autograd-row-wise-of-a-tensor-using-pytorch-autograd-and-without-for-loop/206553/6">Autograd row-wise of a tensor using PyTorch autograd and without for loop</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" style="--category-badge-color: #AB9364; --category-badge-text-color: #FFFFFF;" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    Hi <a class="mention" href="/u/arka_roy">@Arka_Roy</a>, 

So, you’re mixing torch.autograd with torch.func operations, 

which leads to a zero gradient. You need a complete torch.func approach in order for torch.func to return the correct gradient. An example can be found below, 
import torch
from torch.func import jacrev, vmap

device = 'cuda' if torch.cuda.is_available() else 'cpu'

def func(x):
  y1=x**2
  y2=x**3
  return torch.stack([y1,y2],dim=-1)

x = torch.linspace(0,1,10, dtype=torch.float32, device=device) #don't need requires…
  </blockquote>
</aside>