Loss function coming from a PDE

sebnaran: <p>Hello,</p>
<p>I’m trying to train a PINN and I’m computing the physics loss using FEniCS so I am having to detach a torch tensor and then I have to provide an implementation of the gradient. When I train the loss stays constant after each epoch, it doesnt increase or decrease. This probably means that the gradient of the loss that torch is obtaining is zero. I would appreciate some help figuring out what I’m doing wrong as well as any suggestions on the best practice for training PINNs this way.</p>
<p>Here is the implementation that I’m using:</p>
<p>This is the autograd subclass that Im using:</p>
<pre><code class="lang-auto">class PDE_Loss(torch.autograd.Function):
    def __init__(self):
        super(PDE_Loss, self).__init__()
    #p_dofs: this is the output of the neural net
    #A_matrix_params: this is the input to the neural net
    #formulation is a class that computes the necessary vectors and matrices using 
     #FEniCS
    @staticmethod
    def forward(
        ctx,
        p_dofs: torch.tensor,
        A_matrix_params: torch.tensor,
        lossfun: torch.nn.Module,
        formulation: Dg.Darcy_primal_formulation,
        f: torch.tensor,
    ) -&gt; torch.tensor:
        result = lossfun(
            torch.matmul(
                torch.from_numpy(
                    formulation.assemble_linear_system(A_matrix_params.detach().numpy())
                ),
                p_dofs,
            ),
            f,
        )
        ctx.formulation = formulation
        ctx.lossfun = lossfun
        ctx.save_for_backward(
            p_dofs,
            A_matrix_params,
            f,
            result,
        )
        # print(f"{result=}")
        return result
     #Gradient is estimated using finite differences
    @staticmethod
    def backward(ctx, grad_output):
        eps = 1e-8
        (
            p_dofs,
            A_matrix_params,
            f,
            forward_res,
        ) = ctx.saved_tensors

        formulation = ctx.formulation
        lossfun = ctx.lossfun

        grad = []

        for i in range(3):
            A_matrix_params_eps_away = A_matrix_params.clone()
            A_matrix_params_eps_away[i] += eps

            result_eps_away = lossfun(
                torch.matmul(
                    torch.from_numpy(
                        formulation.assemble_linear_system(
                            A_matrix_params_eps_away.numpy()
                        )
                    ),
                    p_dofs,
                ),
                f,
            )
            grad.append(((result_eps_away - forward_res) / eps) * grad_output)
        return None, torch.tensor(np.array(grad)), None, None, None
</code></pre>

KFrank: <p>Hi Sebastian!</p>
<aside class="quote no-group quote-modified" data-username="sebnaran" data-post="1" data-topic="194501" data-full="true">
<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/s/a183cd/48.png" class="avatar"> sebnaran:</div>
<blockquote>
<p>I have to provide an implementation of the gradient.<br>
…<br>
This probably means that the gradient of the loss that torch is obtaining is zero.<br>
…</p>
<pre><code class="lang-auto">     #Gradient is estimated using finite differences
    @staticmethod
    def backward(ctx, grad_output):
        eps = 1e-8
        ...
            A_matrix_params_eps_away[i] += eps
</code></pre>
</blockquote>
</aside>
<p>In principle you should be able to make this approach work.</p>
<p>My first suggestion would be to test your gradient implementation with<br>
a single forward pass / backward pass (rather than observing that your<br>
training doesn’t work and trying to deduce that you’re getting a zero<br>
gradient).  Take a simple test case for which you are comfortable that<br>
you can reliably compute the gradient with some independent method,<br>
say analytically or with an independent numerical computation.</p>
<p>Does your <code>backward()</code> method (as called by running a forward pass<br>
followed by a backward pass) reproduce the known gradient for your<br>
test case?  Does it incorrectly return zero?</p>
<p>Note that your <code>eps = 1e-8</code> is at or below the edge of single-precision<br>
round-off error so that <code>A_matrix_params_eps_away[i] += eps</code> may<br>
end up being numerically equal to <code>A_matrix_params</code>, which would<br>
give you zero for your finite-difference estimate of the gradient.</p>
<p>If you can find an analytic or semi-analytic formula for the gradient, you<br>
will almost certainly be better off using that than numerical differentiation.<br>
If you have to use numerical (finite-difference) differentiation, it should<br>
be possible to make it work adequately well, although such computations<br>
can be nuanced and rather tricky.</p>
<p>Start by making <code>eps</code> bigger than <code>1e-8</code>.  You might consider performing<br>
the numerical differentiation in double precision if doing so is compatible<br>
with your PDE solver.</p>
<p>In general, if you make <code>eps</code> too big, your finite-difference estimate will<br>
become polluted with higher-order terms.  If you make it too small, your<br>
estimate will become noisy due to round-off error, and at some point<br>
<code>eps</code> will, so to speak, round down to zero.  I would suggest that you<br>
take a selection of “typical” use cases and try to empirically determine<br>
a reasonable value for <code>eps</code>.</p>
<p>You might also look for an existing numerical-differentiation package.<br>
A good one will likely give you more robust results than something you<br>
write yourself.  However, such a package will also tend to be a bit tricky<br>
to use – you will probably have to provide it some sort of “control<br>
parameters” that, in effect, give it some guidance as to how large it<br>
should make its internal equivalent of <code>eps</code>.</p>
<p>Good luck!</p>
<p>K. Frank</p>