Revise training process of DoublePendulum

Author: foldfelis

example/DoublePendulum/Project.toml

@@ -2,12 +2,15 @@ name = "DoublePendulum"
 uuid = "0c23c1c1-5f41-4617-a685-ac46aae913c3"
 
 [deps]
+BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+FluxTraining = "7bf95e4d-ca32-48da-9824-f0dc5310474f"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Pluto = "c3e4b0f8-55cb-11ea-2926-15256bba5781"

example/DoublePendulum/src/DoublePendulum.jl

@@ -1,21 +1,87 @@
 module DoublePendulum
 
-using NeuralOperators
-using Flux
-using CUDA
-using JLD2
+using DataDeps, CSV, DataFrames, MLUtils
+using NeuralOperators, Flux
+using CUDA, FluxTraining, BSON
 
-include("data.jl")
+function register_double_pendulum_chaotic()
+    register(DataDep(
+        "DoublePendulumChaotic",
+        """
+        Dataset was generated on the basis of 21 individual runs of a double pendulum.
+        Each of the recorded sequences lasted around 40s and consisted of around 17500 frames.
 
-__init__() = register_double_pendulum_chaotic()
+        * `x_red`: Horizontal pixel coordinate of the red point (the central pivot to the first pendulum)
+        * `y_red`: Vertical pixel coordinate of the red point (the central pivot to the first pendulum)
+        * `x_green`: Horizontal pixel coordinate of the green point (the first pendulum)
+        * `y_green`: Vertical pixel coordinate of the green point (the first pendulum)
+        * `x_blue`: Horizontal pixel coordinate of the blue point (the second pendulum)
+        * `y_blue`: Vertical pixel coordinate of the blue point (the second pendulum)
+
+        Page: https://developer.ibm.com/exchanges/data/all/double-pendulum-chaotic/
+        """,
+        "https://dax-cdn.cdn.appdomain.cloud/dax-double-pendulum-chaotic/2.0.1/double-pendulum-chaotic.tar.gz",
+        "4ca743b4b783094693d313ebedc2e8e53cf29821ee8b20abd99f8fb4c0866f8d",
+        post_fetch_method=unpack
+    ))
+end
+
+function get_data(; i=0, n=-1)
+    data_path = joinpath(datadep"DoublePendulumChaotic", "original", "dpc_dataset_csv")
+    df = CSV.read(
+        joinpath(data_path, "$i.csv"),
+        DataFrame,
+        header=[:x_red, :y_red, :x_green, :y_green, :x_blue, :y_blue]
+    )
+    data = (n < 0) ? collect(Matrix(df)') : collect(Matrix(df)')[:, 1:n]
+
+    return Float32.(data)
+end
+
+function preprocess(𝐱; Δt=1, nx=30, ny=30)
+    # move red point to (0, 0)
+    xs_red, ys_red = 𝐱[1, :], 𝐱[2, :]
+    𝐱[3, :] -= xs_red; 𝐱[5, :] -= xs_red
+    𝐱[4, :] -= ys_red; 𝐱[6, :] -= ys_red
+
+    # needs only green and blue points
+    𝐱 = reshape(𝐱[3:6, 1:Δt:end], 1, 4, :)
+    # velocity of green and blue points
+    ∇𝐱 = 𝐱[:, :, 2:end] - 𝐱[:, :, 1:(end-1)]
+    # merge info of pos and velocity
+    𝐱 = cat(𝐱[:, :, 1:(end-1)], ∇𝐱, dims=1)
+
+    # with info of first nx steps to inference next ny steps
+    n = size(𝐱)[end] - (nx + ny) + 1
+    𝐱s = Array{Float32}(undef, size(𝐱)[1:2]..., nx, n)
+    𝐲s = Array{Float32}(undef, size(𝐱)[1:2]..., ny, n)
+    for i in 1:n
+        𝐱s[:, :, :, i] .= 𝐱[:, :, i:(i+nx-1)]
+        𝐲s[:, :, :, i] .= 𝐱[:, :, (i+nx):(i+nx+ny-1)]
+    end
+
+    return 𝐱s, 𝐲s
+end
 
-function update_model!(model_file_path, model)
-    model = cpu(model)
-    jldsave(model_file_path; model)
-    @warn "model updated!"
+function get_dataloader(; n_file=20, Δt=1, nx=30, ny=30, ratio=0.9, batchsize=100)
+    𝐱s, 𝐲s = Array{Float32}(undef, 2, 4, nx, 0), Array{Float32}(undef, 2, 4, ny, 0)
+    for i in 0:(n_file-1)
+        𝐱s_i, 𝐲s_i = preprocess(get_data(i=i), Δt=Δt, nx=nx, ny=ny)
+        𝐱s, 𝐲s = cat(𝐱s, 𝐱s_i, dims=4), cat(𝐲s, 𝐲s_i, dims=4)
+    end
+
+    data = shuffleobs((𝐱s, 𝐲s))
+    data_train, data_test = splitobs(data, 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(; Δt=1)
+__init__() = register_double_pendulum_chaotic()
+
+function train(; Δt=1, epochs=20)
     if has_cuda()
         @info "CUDA is on"
         device = gpu
@@ -24,46 +90,27 @@ function train(; Δt=1)
         device = cpu
     end
 
-    m = Chain(
-        Dense(2, 64),
-        OperatorKernel(64=>64, (4, 16), FourierTransform, gelu),
-        OperatorKernel(64=>64, (4, 16), FourierTransform, gelu),
-        OperatorKernel(64=>64, (4, 16), FourierTransform, gelu),
-        OperatorKernel(64=>64, (4, 16), FourierTransform),
-        Dense(64, 128, gelu),
-        Dense(128, 2),
-    ) |> device
+    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))
+    loss_func = l₂loss
 
-    loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
+    learner = Learner(
+        model, data, optimiser, loss_func,
+        ToDevice(device, device),
+        Checkpointer(joinpath(@__DIR__, "../model/"))
+    )
 
-    opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
+    fit!(learner, epochs)
 
-    loader_train, loader_test = get_dataloader(Δt=Δt)
-
-    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, 10, leading=false, trailing=true)
-
-    data = [(𝐱, 𝐲) for (𝐱, 𝐲) in loader_train] |> device
-    for e in 1:20
-        @info "Epoch $e\n η: $(opt.os[2].eta)"
-        @time Flux.train!(loss, params(m), data, opt, cb=call_back)
-        (e%3 == 0) && (opt.os[2].eta /= 2)
-    end
+    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

example/DoublePendulum/src/data.jl

@@ -2,5 +2,12 @@ using DoublePendulum
 using Test
 
 @testset "DoublePendulum" begin
-    include("data.jl")
+    xs = DoublePendulum.get_data(i=0, n=100)
+
+    @test size(xs) == (6, 100)
+
+    learner = DoublePendulum.train(epochs=5)
+    loss = learner.cbstate.metricsepoch[ValidationPhase()][:Loss].values[end]
+    @test loss < 0.05
+
 end