Torch.autograd.grad() and backward() call to compute the gradients

Aakhash_Sundaresan: <p>I’m working with Physics Informed Neural Networks and establish the physics by computing the derivatives with the torch.autograd.grad() function. I only calculate the first order derivatives in my framework. After that I try to compute the total loss and call the loss.backward() method on the total loss, which computes the gradients of the loss with respect to all the tensors that have the requires_grad property set to true (involves the weights and biases and also the network inputs with which the network outputs are computed). Am I right?<br>
My question is that what should be the values for options retain_graph and create_graph in the torch.autograd.grad function and how does these values affect the backward() call? If it affects, how does it affect? Can anyone shed light on the same please?</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/aakhash_sundaresan">@Aakhash_Sundaresan</a>,</p>
<p>Do you have a minimal reproducible example for your problem?</p>
<p>If you want to compute higher-order gradients, you need to make sure <code>retain_graph=True</code> and <code>create_graph=True</code> as PyTorch will automatically free the gradients of previous operations (unless you tell it otherwise).</p>

Aakhash_Sundaresan: <pre><code class="lang-auto">def PDE_loss(self, pred, target):
            xy = target
            u, v, p = pred[:, 0:1], pred[:, 1:2], pred[:, 2:3]
            #first order derivative
            u_x = torch.autograd.grad(u, xy, torch.ones([u.shape[0],1]).to(device),retain_graph=True, create_graph=True)[0][:, 0:1]
            u_y = torch.autograd.grad(u, xy, torch.ones([u.shape[0],1]).to(device),retain_graph=True, create_graph=True)[0][:, 1:2]
            v_x = torch.autograd.grad(v, xy, torch.ones([v.shape[0],1]).to(device),retain_graph=True, create_graph=True)[0][:, 0:1]
            v_y = torch.autograd.grad(v, xy, torch.ones([v.shape[0],1]).to(device),retain_graph=True, create_graph=True)[0][:, 1:2]
            p_x = torch.autograd.grad(p, xy, torch.ones([p.shape[0],1]).to(device),retain_graph=True, create_graph=True)[0][:, 0:1]
            p_y = torch.autograd.grad(p, xy, torch.ones([p.shape[0],1]).to(device),retain_graph=True, create_graph=True)[0][:, 1:2]
            #second order derivative
            u_xx = torch.autograd.grad(u_x, xy, torch.ones([u_x.shape[0],1]).to(device),retain_graph=True, create_graph=True)[0][:, 0:1]
            u_yy = torch.autograd.grad(u_y, xy, torch.ones([u_y.shape[0],1]).to(device),retain_graph=True, create_graph=True)[0][:, 1:2]
            v_xx = torch.autograd.grad(v_x, xy, torch.ones([v_x.shape[0],1]).to(device),retain_graph=True, create_graph=True)[0][:, 0:1]
            v_yy = torch.autograd.grad(v_y, xy, torch.ones([v_y.shape[0],1]).to(device),retain_graph=True, create_graph=True)[0][:, 1:2]
            #continous equation
            f0 = u_x + v_y
            #momentum equation
            f1 = u*u_x + v*u_y + p_x - 1/self.Re*(u_xx + u_yy)
            f2 = u*v_x + v*v_y + p_y - 1/self.Re*(v_xx + v_yy)

            mse_f0 = self.MSE_loss(f0, torch.zeros_like(f0))
            mse_f1 = self.MSE_loss(f1, torch.zeros_like(f1))
            mse_f2 = self.MSE_loss(f2, torch.zeros_like(f2))

            return mse_f0 + mse_f1 + mse_f2


    def closure(self):
        total_batch_loss = 0.0
        col_batch_loss = 0.0
        bc_batch_loss = 0.0
        outlet_batch_loss = 0.0

        self.model.train()

        [col_data, bc_data, outlet_data] = self.data[0], self.data[1], self.data[2]
        for i, (col_data, bc_data, outlet_data) in enumerate(zip_longest(col_data, bc_data, outlet_data)):
            total_loss = torch.tensor(0.0).to(device)
            self.optimizer.zero_grad()
            #unpack data
            col_xy = col_data
            bc_xy, bc_uv = bc_data[0], bc_data[1]#
            out_xy, out_p = outlet_data[0], outlet_data[1]

            #forward pass
            if col_xy is not None:
                col_xy = col_xy.clone()
                col_xy.requires_grad = True
                col_pred = self.model(col_xy)
                col_loss = self.PDE_loss(col_pred, col_xy)
                total_loss += col_loss
            if bc_xy is not None:
                bc_pred = self.model(bc_xy)
                bc_loss = self.BC_loss(bc_pred, bc_uv)
                total_loss += bc_loss
            if out_xy is not None:
                out_pred = self.model(out_xy)
                outlet_loss = self.Outlet_loss(out_pred, out_p)
                total_loss += outlet_loss



            total_loss.backward()
            if self.optimizer.__class__.__name__ != "LBFGS":
                self.optimizer.step()
                
            
            total_batch_loss += total_loss.detach().cpu()
            col_batch_loss += col_loss.detach().cpu()
            bc_batch_loss += bc_loss.detach().cpu()
            outlet_batch_loss += outlet_loss.detach().cpu()

            



        self.losses["total"].append(total_batch_loss)
        self.losses["pde"].append(col_batch_loss)
        self.losses["bc"].append(bc_batch_loss)
        self.losses["outlet"].append(outlet_batch_loss)

</code></pre>
<p>In the code above, I don’t see a rationale behind using the retain graph and create graph when calculating the second order gradients. I understand that these options need to be switched on when computing the first order gradients, to be utilized when computing the second order gradients.</p>

AlphaBetaGamma96: <p>When computing second-order gradients, you’ll need make sure the first-order gradient call has <code>retain_graph=True</code> and <code>create_graph=True</code> as you mentioned. You won’t need them for the second-order gradients as you’re not computing any higher-order gradients.</p>
<p>The docs of <code>torch.autograd.grad</code> can be found here: <a href="https://pytorch.org/docs/stable/generated/torch.autograd.grad.html" class="inline-onebox" rel="noopener nofollow ugc">torch.autograd.grad — PyTorch 2.3 documentation</a></p>
<p>When using <code>torch.autograd.grad</code> it will return a Tensor, which represents the gradient of <code>outputs</code> with respect to the <code>inputs</code>. When calling <code>.backward()</code> it will compute gradients in an accumulated fashion on all Tensors that have a <code>.grad</code> attribute. This is why <code>torch.optim.Optim</code> objects, e.g. <code>torch.optim.Adam</code>, require you to call <code>optim.zero_grad()</code> before computing <code>loss.backward()</code>, because if you don’t the <code>.grad</code> attribute of your tensors will have the gradient of the current epoch and the previous epoch.</p>

Aakhash_Sundaresan: <p>Thanks a ton for your timely support. Also, I had this another doubt, why do we require to use the <code>retain_graph</code> in the <code>loss.backward(retain_graph = True)</code> call. What is the rationale behind doing this?</p>

AlphaBetaGamma96: <p>I don’t think you need to have <code>retain_graph=True</code> within <code>loss.backward()</code>, if the interpreter is throwing an error asking you to do <code>loss.backward(retain_graph=True)</code>, perhaps it might be due to mini-batching your gradients? But I’m not 100% sure.</p>

Aakhash_Sundaresan: <p>Suppose, imagine that i compute the gradient of the output of the network using <code>torch.autograd.grad()</code>, with respect to the input, to design the PDE residual loss. And then when I call the <code>loss.backward()</code>, does the gradients get accumulated in the tensors that result from the <code>torch.autograd.grad()</code> call? If so, do I need to use the <code>create_graph = True</code> option in the <code>torch.autograd.grad()</code> call?</p>

AlphaBetaGamma96: <p>If I understand you correctly, you’re asking if the <code>.grad</code> attribute is populated on the PDE residual loss terms (which are calculated via the <code>torch.autograd.grad</code> calls). Would this be the <code>u_xx</code>, <code>u_yy</code> etc. terms?</p>
<p>I don’t think this would be an issue as the only gradient you care about is with respect to the parameters, and you could check if these intermediate have <code>.grad</code> attributes by calling <code>u_xx.grad</code>?</p>
<aside class="quote no-group" data-username="Aakhash_Sundaresan" data-post="7" data-topic="202946">
<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/aakhash_sundaresan/48/66313_2.png" class="avatar"> Aakhash_Sundaresan:</div>
<blockquote>
<p>If so, do I need to use the <code>create_graph = True</code> option in the <code>torch.autograd.grad()</code> call?</p>
</blockquote>
</aside>
<p>You’re performing higher-order derivatives, so I think you should need to keep <code>create_graph=True</code>. If you’re unsure, you can always run your code with <code>create_graph=True</code> and <code>create_graph=False</code> and see if you get the same results.</p>