revise model and save model

Author: foldfelis

example/DoublePendulum/Project.toml

@@ -7,6 +7,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Pluto = "c3e4b0f8-55cb-11ea-2926-15256bba5781"

example/DoublePendulum/model/.gitkeep

@@ -3,12 +3,19 @@ module DoublePendulum
 using NeuralOperators
 using Flux
 using CUDA
+using JLD2
 
 include("data.jl")
 
 __init__() = register_double_pendulum_chaotic()
 
-function train(; loss_bounds=[1, 0.2, 0.1, 0.05, 0.02])
+function update_model!(model_file_path, model)
+    model = cpu(model)
+    jldsave(model_file_path; model)
+    @warn "model updated!"
+end
+
+function train(; loss_bounds=[1, 0.3, 0.1, 0.05])
     if has_cuda()
         @info "CUDA is on"
         device = gpu
@@ -21,6 +28,7 @@ function train(; loss_bounds=[1, 0.2, 0.1, 0.05, 0.02])
         FourierOperator(6=>64, (16, ), relu),
         FourierOperator(64=>64, (16, ), relu),
         FourierOperator(64=>64, (16, ), relu),
+        FourierOperator(64=>64, (16, ), relu),
         FourierOperator(64=>6, (16, )),
     ) |> device
 
@@ -32,20 +40,24 @@ function train(; loss_bounds=[1, 0.2, 0.1, 0.05, 0.02])
 
     data = [(𝐱, 𝐲) for (𝐱, 𝐲) in loader_train] |> device
 
+    losses = Float32[]
     function validate()
         validation_loss = sum(loss(device(𝐱), device(𝐲)) for (𝐱, 𝐲) in loader_test)/length(loader_test)
         @info "loss: $validation_loss"
 
+        push!(losses, validation_loss)
+        (losses[end] == minimum(losses)) && update_model!(joinpath(@__DIR__, "../model/model.jld2"), m)
+
         isempty(loss_bounds) && return
         if validation_loss < loss_bounds[1]
             @warn "change η"
             opt.os[2].eta /= 2
             popfirst!(loss_bounds)
         end
     end
-
-    call_back = Flux.throttle(validate, 1, leading=false, trailing=true)
-    Flux.@epochs 300 @time(Flux.train!(loss, params(m), data, opt, cb=call_back))
+    call_back = Flux.throttle(validate, 10, leading=false, trailing=true)
+    
+    Flux.@epochs 50 @time(Flux.train!(loss, params(m), data, opt, cb=call_back))
 end
 
 end