Error when a neural ode is implemented

Ashima_Kalathingal · September 27, 2024, 1:33pm

I am trying to implement a neural ode in Julia. The loss function reduces when training. But after the training the parameters are not updated. It goes back to initial value. When I run the algorithm the following error shows. I don’t know whether this happens because of it

Warning: Potential performance improvement omitted. EnzymeVJP tried and failed in the automated AD choice algorithm. To show the stack trace, set SciMLSensitivity.STACKTRACE_WITH_VJPWARN[] = true. To turn off this printing, add `verbose = false` to the `solve` call.
└ @ SciMLSensitivity C:\Users\Kalath_A\.julia\packages\SciMLSensitivity\PstNN\src\concrete_solve.jl:24

┌ Warning: Lux.apply(m::AbstractExplicitLayer, x::AbstractArray{<:ReverseDiff.TrackedReal}, ps, st) input was corrected to Lux.apply(m::AbstractExplicitLayer, x::ReverseDiff.TrackedArray}, ps, st).
│ 
│ 1. If this was not the desired behavior overload the dispatch on `m`.
│ 
│ 2. This might have performance implications. Check which layer was causing this problem using `Lux.Experimental.@debug_mode`.

Is this error affecting my code? What does this error mean? Can anyone help me

avikpal · September 27, 2024, 8:48pm

Can you share your code? If loss is going down, it is very likely that you are not using inplace operations and the parameters you get returned are not the trained parameters.

The 2 warnings (not an error) are saying:

could not use Enzyme.jl for AD
Lux.jl detected a AbstractArray of TrackedReals from ReverseDiff instead of TrackedArray. Since the former almost certainly doesn’t have any benefits, it is converting it to a TrackedArray

Ashima_Kalathingal · September 29, 2024, 8:46pm

Thank you for your reply.
The complete code is below.

import Pkg
Pkg.activate("D:/ASHIMA/JULIA/Neuralode")

# Importing packages
using Plots, JLD2, Statistics, ComponentArrays,Random  
using DifferentialEquations, Lux   
using Optimization, SciMLSensitivity, OptimizationFlux,OptimizationOptimJL,Zygote
using StableRNGs
using Interpolations

function find_discharge_end(Current_data,start=5)
    for i in start:length(Current_data)
        if abs(Current_data[i]) == 0
            return i 
        end
    end
    return -1 
end

function current_val(Crate)
    if Crate == "0p5C"
        return 0.5*5.0
    elseif Crate == "1C"
        return 1.0*5.0
    elseif Crate == "2C"
        return 2.0*5.0
    elseif Crate == "1p5C"
        return 1.5*5.0
    end
end

# Function to load the training data
data_file = load("Datasets_ashima_WLTP.jld2")["Datasets_WLTP"]

function load_data(name1::String,name2::Int64)
    if name1 == "WLTP"
        data = data_file["$(name1)_T$(name2)"]
        I    = LinearInterpolation(data["time"],data["current"],extrapolation_bc=Interpolations.Flat())
        t,T,T∞ = data["time"],data["temperature"],data["temperature"][1]
        SOC0   = 0.9
        len    = length(t)
        return t,T,T∞,I,SOC0,len
    else
        data = data_file["$(name1)_T$(name2)"]
        n    = find_discharge_end(data["current"])
        t,T,T∞ = data["time"][2:n],data["temperature"][2:n],data["temperature"][1]
        constant_current = current_val(name1)
        I = t -> constant_current
        SOC0   = 1.0
        len    = length(t)
        return t,T,T∞,I,SOC0,len
    end
end

#training conditions     
Crate1,Temp1 = "1C",10
Crate2,Temp2 = "0p5C",25
Crate3,Temp3 = "WLTP",0
Crate4,Temp4 = "1C",25
Crate5,Temp5 = "WLTP",25
Crate6,Temp6 = "2C",10

t1,T1,T∞1,I1,SOC1_0,len1     = load_data(Crate1,Temp1)
t2,T2,T∞2,I2,SOC2_0,len2     = load_data(Crate2,Temp2)
t3,T3,T∞3,I3,SOC3_0,len3     = load_data(Crate3,Temp3)
t4,T4,T∞4,I4,SOC4_0,len4     = load_data(Crate4,Temp4)
t5,T5,T∞5,I5,SOC5_0,len5     = load_data(Crate5,Temp5)
t6,T6,T∞6,I6,SOC6_0,len6     = load_data(Crate6,Temp6)

# 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)

function UDE_model!(du,u,p,t,T∞,I_func)

    Q = 4.29*3600 #C/2 Cell capacity in As
    C₁, C₂ = 0.0015397895, 0.020306583        
    
    I = I_func(t) # Function of current
    du[1] = -I/Q  #SOC update

    G= (NN([u[1],u[2],I],p, _st_)[1][1])^2

    du[2] = -C₁ * (u[2] - T∞) + C₂ *G
    

    return du

end 

UDE_model1!(du,u,p,t) = UDE_model!(du,u,p,t,T∞1,I1)
UDE_model2!(du,u,p,t) = UDE_model!(du,u,p,t,T∞2,I2)
UDE_model3!(du,u,p,t) = UDE_model!(du,u,p,t,T∞3,I3)
UDE_model4!(du,u,p,t) = UDE_model!(du,u,p,t,T∞4,I4)
UDE_model5!(du,u,p,t) = UDE_model!(du,u,p,t,T∞5,I5)
UDE_model6!(du,u,p,t) = UDE_model!(du,u,p,t,T∞6,I6)

prob1 = ODEProblem(UDE_model1!,[SOC1_0,T∞1],(t1[1],t1[end]),Para)
prob2 = ODEProblem(UDE_model2!,[SOC2_0,T∞2],(t2[1],t2[end]),Para)
prob3 = ODEProblem(UDE_model3!,[SOC3_0,T∞3],(t3[1],t3[end]),Para)
prob4 = ODEProblem(UDE_model4!,[SOC4_0,T∞4],(t4[1],t4[end]),Para)
prob5 = ODEProblem(UDE_model5!,[SOC5_0,T∞5],(t5[1],t5[end]),Para)
prob6 = ODEProblem(UDE_model6!,[SOC6_0,T∞6],(t6[1],t6[end]),Para)

solver=Tsit5()
sol1 = solve(prob1, solver, saveat = t1)
sol2 = solve(prob2, solver, saveat = t2)
sol3 = solve(prob3, solver, saveat = t3)
sol4 = solve(prob4, solver, saveat = t4)
sol5 = solve(prob5, solver, saveat = t5)
sol6 = solve(prob6, solver, saveat = t6)

😄 = 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) 
optprob   = Optimization.OptimizationProblem(optf, Para) 
result    = Optimization.solve(optprob, optimiser, callback = callback, maxiters = 510)
Para_opt_ADAM      = result.u

If you find any error please let me know. Any help would be much appreciated.