How to find partial derivative in pytorch

Question

I have a model u(x,t) with layers 2X50 , then 50X50 , and 50X1 .

I train the model with input x,t of size [100,2] . In the final layer I get u and now I want to differentiate it w.r.t to x and t and double differentiate w.r.t to x . How do I do this in PyTorch?

Answer 1

You can use PyTorch's autograd engine like so:

import torch

x = torch.randn(100, requires_grad=True)
t = torch.randn(2, requires_grad=True)
u = u(x,t)

# 1st derivatives
dt = torch.autograd.grad(u, t)[0]
dx = torch.autograd.grad(u, x, create_graph=True)[0]

# 2nd derivatives (higher orders require `create_graph=True`)
ddx = torch.autograd.grad(dx, x)[0]