Hi All, I am relatively new to Julia, coming from a C/C++ & Python background. I am absolutely loving Julia, but I am struggling a bit to build a simple PINN model using Zygote/Flux (and I tried ForwardDiff). I am using one of the standard examples from Raissi et al. Simple PINN NLS, but I simplified the loss function for testing purposes. I would like to define a loss function which has derivatives with respect to input features, and then use Zygote.gradient to find the gradients of the loss function with respect to the model parameters.
Please see below the definition of the network:
#define layers fc1 = Dense(2, 20, tanh); fc2 = Dense(20, 2, identity); #create chain and set optimizer model = Chain(fc1, fc2); pars = Flux.params(model); opt = ADAM();
Here is my loss function code:
function loss_fn(model, X) u(x) = model(x)[1,:] v(x) = model(x)[2,:] ∂u∂x = Zygote.gradient(a -> sum(u(a)), X)[1,:] ∂u∂t = Zygote.gradient(a -> sum(u(a)), X)[2,:] ∂v∂x = Zygote.gradient(a -> sum(v(a)), X)[1,:] ∂v∂t = Zygote.gradient(a -> sum(v(a)), X)[2,:] loss = mean(u(X).^2) + mean(v(X).^2) + mean(∂u∂x.^2) + mean(∂u∂t.^2) + mean(∂v∂x.^2) + mean(∂v∂t.^2) return loss end
Here is the line I am trying to use to find the gradients of the loss function to model parameters:
grads = Flux.gradient(() -> loss_fn(model, x0_t0), pars)
But when I run the above code I get a “mutating” array error. I know other colleagues have proposed using a combination of ReverseDiff and Zygote.gradient to fix this. But honestly I am not sure what that would look like? Any help would be appreciated.