I’m trying to implement a Universal Differential Equation (UDE) based on a 2D diffusion PDE with an integrated neural network. I’m computing the diffusivity of a 2D matrix H with a forward scheme, and once I reach the the final state of H I compute the loss function and try to backpropagate the whole thing using Zygote.pullback() in order to update the weights of the neural network.
Since I am updating (overwriting) the H matrix at each time step, Zygote rightfully complains that “Mutating arrays is not supported”. I have tried to work around this by using Zygote.Buffer, but I still keep getting the error. I have also read most posts related to this issue, but I haven’t been able to find a solution to this. I’m using Julia 1.6.
Here is an overview of what my code does (I have omitted many parts, as the original code does many more things that I judge irrelevant to this issue):
# Train the UDE
# We get the model parameters to be trained
ps_UA = Flux.params(UA) # get the parameters of the NN
#data = Flux.Data.DataLoader((X, Y), batchsize=hyparams.batchsize, (X, Y), shuffle=false)
data = Dict("X"=>X, "Y"=>Y) # Get the input data
# Train the UDE for a given number of epochs
@epochs hyparams.epochs hybrid_train!(loss, ps_UA, data, opt, H, p, t, t₁)
Which calls
function hybrid_train!(loss, ps_UA, data, opt, H, p, t, t₁)
# Some previous code here
# Retrieve model parameters
ps_UA = Params(ps_UA)
# back is a method that computes the product of the gradient so far with its argument.
train_loss_UA, back_UA = Zygote.pullback(() -> loss(data, H, p, t, t₁), ps_UA)
# Callback to track the training
callback(train_loss_UA)
# Apply back() to the correct type of 1.0 to get the gradient of loss.
gs_UA = back_UA(one(train_loss_UA))
# Insert what ever code you want here that needs gradient.
Flux.update!(opt, ps_UA, gs_UA)
end
Based on a loss function which contains the forward scheme
function loss(data, H, p, t, t₁)
l_H = 0.0f0
# Compute l_H as the difference between the simulated H with UA(x) and H_ref
iceflow!(H, UA, p,t,t₁)
l_H = sqrt(Flux.Losses.mse(H, H_ref["H"][end]; agg=mean))
return l_H
Where the forward scheme is defined here (this is a simplified version of the code):
function iceflow!(H, UA, p,t,t₁)
# Manual explicit forward scheme implementation
while t < t₁
# Long chunk of code which computes the flux F in a staggered grid
# Δt is defined to be stable
dHdt = (F .+ inn(MB[:,:,year])) .* Δt # dH for time step t
# Use Zygote.Buffer in order to mutate matrix while being able to compute gradients
H_buff = Zygote.Buffer(H)
H_buff[2:end - 1,2:end - 1] .= max.(0.0, inn(H_buff) .+ dHdt) # CODE FAILS HERE
H = copy(H_Zygote)
t += Δt
end
end
So my problem is that I’m overwriting the H matrix at each time step, but Zygote cannot perform the pullback on that. Is there any clean workaround to this? Zygote.Buffer doesn’t seem to help here. Thanks in advance!