Refactor

Author: foldfelis

example/Burgers/src/Burgers.jl

@@ -44,18 +44,19 @@ end
 
 __init__() = register_burgers()
 
-function train(; epochs=500)
-    if has_cuda()
-        @info "CUDA is on"
+function train(; cuda=true, η₀=1f-3, λ=1f-4, epochs=500)
+    if cuda && CUDA.has_cuda()
         device = gpu
         CUDA.allowscalar(false)
+        @info "Training on GPU"
     else
         device = cpu
+        @info "Training on CPU"
     end
 
     model = FourierNeuralOperator(ch=(2, 64, 64, 64, 64, 64, 128, 1), modes=(16, ), σ=gelu)
     data = get_dataloader()
-    optimiser = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
+    optimiser = Flux.Optimiser(WeightDecay(λ), Flux.ADAM(η₀))
     loss_func = l₂loss
 
     learner = Learner(

example/DoublePendulum/src/DoublePendulum.jl

@@ -81,18 +81,19 @@ end
 
 __init__() = register_double_pendulum_chaotic()
 
-function train(; Δt=1, epochs=20)
-    if has_cuda()
-        @info "CUDA is on"
+function train(; cuda=true, Δt=1, η₀=1f-3, λ=1f-4, epochs=20)
+    if cuda && CUDA.has_cuda()
         device = gpu
         CUDA.allowscalar(false)
+        @info "Training on GPU"
     else
         device = cpu
+        @info "Training on CPU"
     end
 
     model = FourierNeuralOperator(ch=(2, 64, 64, 64, 64, 64, 128, 2), modes=(4, 16), σ=gelu)
     data = get_dataloader(Δt=Δt)
-    optimiser = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
+    optimiser = Flux.Optimiser(WeightDecay(λ), Flux.ADAM(η₀))
     loss_func = l₂loss
 
     learner = Learner(

example/FlowOverCircle/src/FlowOverCircle.jl

@@ -48,18 +48,19 @@ function get_dataloader(; ts::AbstractRange=LinRange(100, 11000, 10000), ratio::
     return loader_train, loader_test
 end
 
-function train(; epochs=50)
-    if has_cuda()
-        @info "CUDA is on"
+function train(; cuda=true, η₀=1f-3, λ=1f-4, epochs=50)
+    if cuda && CUDA.has_cuda()
         device = gpu
         CUDA.allowscalar(false)
+        @info "Training on GPU"
     else
         device = cpu
+        @info "Training on CPU"
     end
 
     model = MarkovNeuralOperator(ch=(1, 64, 64, 64, 64, 64, 1), modes=(24, 24), σ=gelu)
     data = get_dataloader()
-    optimiser = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
+    optimiser = Flux.Optimiser(WeightDecay(λ), Flux.ADAM(η₀))
     loss_func = l₂loss
 
     learner = Learner(
@@ -73,17 +74,17 @@ function train(; epochs=50)
     return learner
 end
 
-function train_gno(; epochs=50)
-    if has_cuda()
-        @info "CUDA is on"
+function train_gno(; cuda=true, η₀=1f-3, λ=1f-4, epochs=50)
+    if cuda && CUDA.has_cuda()
         device = gpu
         CUDA.allowscalar(false)
+        @info "Training on GPU"
     else
         device = cpu
+        @info "Training on CPU"
     end
 
     featured_graph = FeaturedGraph(grid([96, 64]))
-
     model = Chain(
         Dense(1, 16),
         WithGraph(featured_graph, GraphKernel(Dense(2*16, 16, gelu), 16)),
@@ -93,7 +94,7 @@ function train_gno(; epochs=50)
         Dense(16, 1),
     )
     data = get_dataloader(batchsize=16, flatten=true)
-    optimiser = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
+    optimiser = Flux.Optimiser(WeightDecay(λ), Flux.ADAM(η₀))
     loss_func = l₂loss
 
     learner = Learner(

example/SuperResolution/src/SuperResolution.jl

@@ -69,18 +69,19 @@ function fit!(learner, nepochs::Int)
     fit!(learner, nepochs, (learner.data.training, learner.data.validation, learner.data.testing))
 end
 
-function train(; epochs=50)
-    if has_cuda()
-        @info "CUDA is on"
+function train(; cuda=true, η₀=1f-3, λ=1f-4, epochs=50)
+    if cuda && CUDA.has_cuda()
         device = gpu
         CUDA.allowscalar(false)
+        @info "Training on GPU"
     else
         device = cpu
+        @info "Training on CPU"
     end
 
     model = MarkovNeuralOperator(ch=(1, 64, 64, 64, 64, 64, 1), modes=(24, 24), σ=gelu)
     data = get_dataloader()
-    optimiser = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
+    optimiser = Flux.Optimiser(WeightDecay(λ), Flux.ADAM(η₀))
     loss_func = l₂loss
 
     learner = Learner(
@@ -101,34 +102,4 @@ function get_model()
     return BSON.load(joinpath(model_path, model_file), @__MODULE__)[:model]
 end
 
-# using NeuralOperators
-# using Flux
-# using Flux.Losses: mse
-# using Flux.Data: DataLoader
-# using GeometricFlux
-# using Graphs
-# using CUDA
-# using JLD2
-# using ProgressMeter: Progress, next!
-
-# include("data.jl")
-# include("models.jl")
-
-# function update_model!(model_file_path, model)
-#     model = cpu(model)
-#     jldsave(model_file_path; model)
-#     @info "model updated!"
-# end
-
-# function get_model()
-#     f = jldopen(joinpath(@__DIR__, "../model/model.jld2"))
-#     model = f["model"]
-#     close(f)
-
-#     return model
-# end
-
-# loss(m, 𝐱, 𝐲) = mse(m(𝐱), 𝐲)
-# loss(m, loader::DataLoader, device) = sum(loss(m, 𝐱 |> device, 𝐲 |> device) for (𝐱, 𝐲) in loader)/length(loader)
-
 end # module