I am implementing a Neural ODE in julia. When I complete the training using ADAM optimizer, the loss function decreases as expected. But after training, the parameters reset to the original initial value. So I am not able to obtain the updated parameters after training.
Has anyone experienced a similar issue?
I will paste some of my code here.
This is how I initialize the neural network,
# Neural network
NN = Lux.Chain(Lux.Dense(3,20,tanh),Lux.Dense(20,20,tanh),Lux.Dense(20,1))
rng = StableRNG(11)
Para0,st = Lux.setup(rng,NN) #Initializing parameter and state of the neural network
const _st_ = st # A constant parameter _st_ is created with the value of st
Para = ComponentVector(Para0)
# The UDE_model! returns the derivative of the problem
UDE_model1!(du,u,p,t) = UDE_model!(du,u,p,t,Tβ1,I1)
# There are 6 such cases
prob1 = ODEProblem(UDE_model1!,[SOC1_0,Tβ1],(t1[1],t1[end]),Para)
#There are 6 such problems
solver=Tsit5()
sol1 = solve(prob1, solver, saveat = t1)
#There are 6 such solutions
# The loss function is defined as below
π = 2
function loss_UDE6(ΞΈ)
N_dataset = 6
Solver = Tsit5()
if π% N_dataset == 0
_prob = remake(prob1, p = ΞΈ)
_sol = Array(solve(_prob, Solver, saveat = t1))
e1 = sum(abs2, T1 .- _sol[2, :]) / len1
println("Loss for $(Crate1) $(Temp1) is $(sqrt(e1))")
return e1
elseif π% N_dataset == 1
_prob = remake(prob2, p = ΞΈ)
_sol = Array(solve(_prob, Solver, saveat = t2))
e2 = sum(abs2, T2 .- _sol[2, :]) / len2
println("Loss for $(Crate2) $(Temp2) is $(sqrt(e2))")
return e2
elseif π% N_dataset == 2
_prob = remake(prob3, p = ΞΈ)
_sol = Array(solve(_prob, Solver, saveat = t3))
e3 = sum(abs2, T3 .- _sol[2, :]) / len3
println("Loss for $(Crate3) $(Temp3) is $(sqrt(e3))")
return e3
elseif π% N_dataset == 3
_prob = remake(prob4, p = ΞΈ)
_sol = Array(solve(_prob, Solver, saveat = t4))
e4 = sum(abs2, T4 .- _sol[2, :]) / len4
println("Loss for $(Crate4) $(Temp4) is $(sqrt(e4))")
return e4
elseif π% N_dataset == 4
_prob = remake(prob5, p = ΞΈ)
_sol = Array(solve(_prob, Solver, saveat = t5))
e5 = sum(abs2, T5 .- _sol[2, :]) / len5
println("Loss for $(Crate5) $(Temp5) is $(sqrt(e5))")
return e5
elseif π% N_dataset == 5
_prob = remake(prob6, p = ΞΈ)
_sol = Array(solve(_prob, Solver, saveat = t6))
e6 = sum(abs2, T6 .- _sol[2, :]) / len6
println("Loss for $(Crate6) $(Temp6) is $(sqrt(e6))")
return e6
end
end
Itera = 1
plot_ = plot(framestyle = :box, legend = :none, xlabel = "Iteration", ylabel = "Loss function", title = "Training neural network")
function callback(p,l)
global π += 1
global Itera += 1
colors_ = [:red, :green, :blue, :orange, :purple, :brown]
println("Objective value at iteration $(Itera) is $(sqrt(l)) ")
scatter!(plot_, [Itera], [sqrt(l)], markersize = 4, markercolor = colors_[π % 6 + 1])
display(plot_)
return false
end
optimiser = ADAM(3e-4)
AD_type = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x,p) -> loss_UDE6(x), AD_type) # Defines the optimization function. x represents the parameters to be optimized. p is used for extra parameters but not used in this setup
optprob = Optimization.OptimizationProblem(optf, Para) # State the optimization function and the initial value of parameters
result = Optimization.solve(optprob, optimiser, callback = callback, maxiters = 510)
Para_opt_ADAM = result.u
After I run this code the Para_opt_ADAM still gives the initial value of the weights and biases. The loss function is decreasing as the training progresses. But somehow after the training progresses the result is not saved.
Can someone help me on this? I donβt understand where I am going wrong.