Revise training process for FlowOverCircle

Author: foldfelis

.gitignore

@@ -25,4 +25,4 @@ docs/site/
 # environment.
 Manifest.toml
 
-*.jld2
+*.bson

example/FlowOverCircle/Project.toml

@@ -2,13 +2,16 @@ name = "FlowOverCircle"
 uuid = "1fc04e5d-1dd1-42ff-8d75-1d53504b2476"
 
 [deps]
+BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
-JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+FluxTraining = "7bf95e4d-ca32-48da-9824-f0dc5310474f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Pluto = "c3e4b0f8-55cb-11ea-2926-15256bba5781"
+ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 WaterLily = "ed894a53-35f9-47f1-b17f-85db9237eebd"
 
 [extras]

example/FlowOverCircle/src/FlowOverCircle.jl

@@ -1,16 +1,44 @@
 module FlowOverCircle
 
-using NeuralOperators
-using Flux
-using CUDA
-using JLD2
+using WaterLily, LinearAlgebra, ProgressMeter, MLUtils
+using NeuralOperators, Flux
+using CUDA, FluxTraining, BSON
 
-include("data.jl")
+function circle(n, m; Re=250) # copy from [WaterLily](https://github.com/weymouth/WaterLily.jl)
+    # Set physical parameters
+    U, R, center = 1., m/8., [m/2, m/2]
+    ν = U * R / Re
 
-function update_model!(model_file_path, model)
-    model = cpu(model)
-    jldsave(model_file_path; model)
-    @warn "model updated!"
+    body = AutoBody((x,t) -> LinearAlgebra.norm2(x .- center) - R)
+    Simulation((n+2, m+2), [U, 0.], R; ν, body)
+end
+
+function gen_data(ts::AbstractRange)
+    @info "gen data... "
+    p = Progress(length(ts))
+
+    n, m = 3(2^5), 2^6
+    circ = circle(n, m)
+
+    𝐩s = Array{Float32}(undef, 1, n, m, length(ts))
+    for (i, t) in enumerate(ts)
+        sim_step!(circ, t)
+        𝐩s[1, :, :, i] .= Float32.(circ.flow.p)[2:end-1, 2:end-1]
+
+        next!(p)
+    end
+
+    return 𝐩s
+end
+
+function get_dataloader(; ts::AbstractRange=LinRange(100, 11000, 10000), ratio::Float64=0.95, batchsize=100)
+    data = gen_data(ts)
+    data_train, data_test = splitobs((𝐱=data[:, :, :, 1:end-1], 𝐲=data[:, :, :, 2:end]), at=ratio)
+
+    loader_train = Flux.DataLoader(data_train, batchsize=batchsize, shuffle=true)
+    loader_test = Flux.DataLoader(data_test, batchsize=batchsize, shuffle=false)
+
+    return loader_train, loader_test
 end
 
 function train()
@@ -22,42 +50,34 @@ function train()
         device = cpu
     end
 
-    m = Chain(
+    model = Chain(
         Dense(1, 64),
         OperatorKernel(64=>64, (24, 24), FourierTransform, gelu),
         OperatorKernel(64=>64, (24, 24), FourierTransform, gelu),
         OperatorKernel(64=>64, (24, 24), FourierTransform, gelu),
         OperatorKernel(64=>64, (24, 24), FourierTransform, gelu),
         Dense(64, 1),
-    ) |> device
+    )
+    data = get_dataloader()
+    optimiser = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
+    loss_func = l₂loss
 
-    loss(𝐱, 𝐲) = l₂loss(m(𝐱), 𝐲)
+    learner = Learner(
+        model, data, optimiser, loss_func,
+        ToDevice(device, device),
+        Checkpointer(joinpath(@__DIR__, "../model/"))
+    )
 
-    opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
-
-    @info "gen data... "
-    @time loader_train, loader_test = get_dataloader()
+    fit!(learner, 50)
 
-    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)
-    end
-    call_back = Flux.throttle(validate, 5, leading=false, trailing=true)
-
-    data = [(𝐱, 𝐲) for (𝐱, 𝐲) in loader_train] |> device
-    Flux.@epochs 50 @time(Flux.train!(loss, params(m), data, opt, cb=call_back))
+    return learner
 end
 
 function get_model()
-    f = jldopen(joinpath(@__DIR__, "../model/model.jld2"))
-    model = f["model"]
-    close(f)
+    model_path = joinpath(@__DIR__, "../model/")
+    model_file = readdir(model_path)[end]
 
-    return model
+    return BSON.load(joinpath(model_path, model_file), @__MODULE__)[:model]
 end
 
-end
+end # module