revise model

Author: foldfelis

example/DoublePendulum/src/DoublePendulum.jl

@@ -18,25 +18,36 @@ function train()
     end
 
     m = Chain(
-        FourierOperator(6=>6, (16, ), gelu),
-        FourierOperator(6=>6, (16, ), gelu),
-        FourierOperator(6=>6, (16, ), gelu),
+        FourierOperator(6=>6, (16, ), relu),
+        FourierOperator(6=>6, (16, ), relu),
+        FourierOperator(6=>6, (16, ), relu),
+        FourierOperator(6=>6, (16, ), relu),
         FourierOperator(6=>6, (16, )),
     ) |> device
 
     loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
 
+    opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
+
     loader_train, loader_test = get_dataloader()
 
+    data = [(𝐱, 𝐲) for (𝐱, 𝐲) in loader_train] |> device
+
+    loss_bounds = [0.3, 0.05, 0.01]
     function validate()
-        validation_losses = [loss(device(𝐱), device(𝐲)) for (𝐱, 𝐲) in loader_test]
-        @info "loss: $(sum(validation_losses)/length(loader_test))"
+        validation_loss = sum(loss(device(𝐱), device(𝐲)) for (𝐱, 𝐲) in loader_test)/length(loader_test)
+        @info "loss: $validation_loss"
+
+        isempty(loss_bounds) && return
+        if validation_loss < loss_bounds[1]
+            @warn "change η"
+            opt.os[2].eta /= 2
+            popfirst!(loss_bounds)
+        end
     end
 
-    data = [(𝐱, 𝐲) for (𝐱, 𝐲) in loader_train] |> device
-    opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-4))
-    call_back = Flux.throttle(validate, 0.5, leading=false, trailing=true)
-    Flux.@epochs 500 @time(Flux.train!(loss, params(m), data, opt, cb=call_back))
+    call_back = Flux.throttle(validate, 1, leading=false, trailing=true)
+    Flux.@epochs 300 @time(Flux.train!(loss, params(m), data, opt, cb=call_back))
 end
 
 end