Saving and loading model and the best weight after training in SciML Julia

user_12976585 0
Viewed 206

I am trying to understand how could we save the model + its best weight after training it using sciml_train within the the Neural ODE context.

After:

result_neuralode = DiffEqFlux.sciml_train(loss_neuralode, prob_neuralode.p,
                                          ADAM(0.05), cb = callback,
                                          maxiters = 100)

I don't really understand how do we save the model + best weight. I know that the best weight can be accessed by calling result_neuralode.minimizer. I am looking for an equivalent method to tensorflow.save_model(fn="mymodel").

I have tried to use this:

result_neuralode = DiffEqFlux.sciml_train(loss_neuralode, prob_neuralode.p,
                                          ADAM(0.05), cb = callback,
                                          maxiters = 100)

println("The  best weight is: ")
display(result_neuralode.minimizer)

#Access the weight of the model
weights = params(prob_neuralode)

#save weights to fn: "weight.bson"
using BSON: @save
@save "weight.bson" weights

#load fn to variable called W
using BSON: @load
@load "weight.bson" W

#mounting the loaded weight (W)
Flux.loadparams!(prob_neuralode, W)

#predict using the loaded weight
data = predict_neuralode(W)

#plot the loaded prediction
plot!(
    tsteps,data[1,:],label="Loaded"
)

It throws me error:

ERROR: LoadError: UndefVarError: Zygote not defined

I'm a bit confused on how can I save the model + its weight, and reload it sometimes later for an actual application.

julia
1 Answers

You have to do using Zygote before loading IIRC. But if you save the array result_neuralode.u then you have all of the weights and don't need to save anything else to reconstruct the system.

user_15442030
Related