在 SciML Julia 中训练后保存和加载模型和最佳权重

Question

我试图了解如何在神经 ODE 上下文中使用 sciml_train 训练模型后保存模型 + 其最佳权重。

之后：

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

我真的不明白我们如何保存模型 + 最佳权重。我知道可以通过调用 result_neuralode.minimizer 获得最佳权重。我正在寻找与 tensorflow.save_model(fn="mymodel") 等效的方法。

我曾尝试使用这个：

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

它抛出了我的错误：

ERROR: LoadError: UndefVarError: Zygote not defined

我对如何保存模型 + 它的重量感到有些困惑，并且有时稍后为实际应用重新加载它。

Answer 1

您必须在加载 IIRC 之前执行 using Zygote。但是如果你保存数组 result_neuralode.u，那么你就拥有了所有的权重，不需要保存任何其他东西来重建系统。