Merge pull request #57 from SciML/training_process

Revise training process

Author: foldfelis

.gitignore

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

Project.toml

@@ -16,9 +16,9 @@ Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-CUDA = "3.8"
+CUDA = "3.9"
 CUDAKernels = "0.3, 0.4"
-ChainRulesCore = "1.13"
+ChainRulesCore = "1.14"
 FFTW = "1.4"
 Flux = "0.13"
 GeometricFlux = "0.11"

example/Burgers/Project.toml

@@ -2,10 +2,13 @@ name = "Burgers"
 uuid = "5b053d85-f964-4905-ae31-99551cd8d3ad"
 
 [deps]
+BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+FluxTraining = "7bf95e4d-ca32-48da-9824-f0dc5310474f"
 MAT = "23992714-dd62-5051-b70f-ba57cb901cac"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 
 [extras]

example/Burgers/src/Burgers.jl

@@ -1,51 +1,72 @@
 module Burgers
 
-using NeuralOperators
-using Flux
-using CUDA
+using DataDeps, MAT, MLUtils
+using NeuralOperators, Flux
+using CUDA, FluxTraining, BSON
 
-include("data.jl")
 include("Burgers_deeponet.jl")
 
+function register_burgers()
+    register(DataDep(
+        "Burgers",
+        """
+        Burgers' equation dataset from
+        [fourier_neural_operator](https://github.com/zongyi-li/fourier_neural_operator)
+        """,
+        "http://www.med.cgu.edu.tw/NeuralOperators/Burgers_R10.zip",
+        "9cbbe5070556c777b1ba3bacd49da5c36ea8ed138ba51b6ee76a24b971066ecd",
+        post_fetch_method=unpack
+    ))
+end
+
+function get_data(; n=2048, Δsamples=2^3, grid_size=div(2^13, Δsamples), T=Float32)
+    file = matopen(joinpath(datadep"Burgers", "burgers_data_R10.mat"))
+    x_data = T.(collect(read(file, "a")[1:n, 1:Δsamples:end]'))
+    y_data = T.(collect(read(file, "u")[1:n, 1:Δsamples:end]'))
+    close(file)
+
+    x_loc_data = Array{T, 3}(undef, 2, grid_size, n)
+    x_loc_data[1, :, :] .= reshape(repeat(LinRange(0, 1, grid_size), n), (grid_size, n))
+    x_loc_data[2, :, :] .= x_data
+
+    return x_loc_data, reshape(y_data, 1, :, n)
+end
+
+function get_dataloader(; ratio::Float64=0.9, batchsize=100)
+    𝐱, 𝐲 = get_data(n=2048)
+    data_train, data_test = splitobs((𝐱, 𝐲), at=ratio)
+
+    loader_train = DataLoader(data_train, batchsize=batchsize, shuffle=true)
+    loader_test = DataLoader(data_test, batchsize=batchsize, shuffle=false)
+
+    return loader_train, loader_test
+end
+
 __init__() = register_burgers()
 
-function train()
-    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
 
-    modes = (16, )
-    ch = 64 => 64
-    σ = gelu
-    Transform = FourierTransform
-    m = Chain(
-        Dense(2, 64),
-        OperatorKernel(ch, modes, Transform, σ),
-        OperatorKernel(ch, modes, Transform, σ),
-        OperatorKernel(ch, modes, Transform, σ),
-        OperatorKernel(ch, modes, Transform),
-        Dense(64, 128, σ),
-        Dense(128, 1),
-        flatten
-    ) |> device
-
-    loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
-
-    loader_train, loader_test = get_dataloader()
-
-    function validate()
-        validation_losses = [loss(device(𝐱), device(𝐲)) for (𝐱, 𝐲) in loader_test]
-        @info "loss: $(sum(validation_losses)/length(loader_test))"
-    end
+    model = FourierNeuralOperator(ch=(2, 64, 64, 64, 64, 64, 128, 1), modes=(16, ), σ=gelu)
+    data = get_dataloader()
+    optimiser = Flux.Optimiser(WeightDecay(λ), Flux.ADAM(η₀))
+    loss_func = l₂loss
+
+    learner = Learner(
+        model, data, optimiser, loss_func,
+        ToDevice(device, device),
+    )
+
+    fit!(learner, epochs)
 
-    data = [(𝐱, 𝐲) for (𝐱, 𝐲) in loader_train] |> device
-    opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
-    call_back = Flux.throttle(validate, 5, leading=false, trailing=true)
-    Flux.@epochs 500 @time(Flux.train!(loss, params(m), data, opt, cb=call_back))
+    return learner
 end
 
 end

example/Burgers/src/Burgers_deeponet.jl

@@ -1,3 +1,12 @@
+function get_data_don(; n=2048, Δsamples=2^3, grid_size=div(2^13, Δsamples))
+    file = matopen(joinpath(datadep"Burgers", "burgers_data_R10.mat"))
+    x_data = collect(read(file, "a")[1:n, 1:Δsamples:end])
+    y_data = collect(read(file, "u")[1:n, 1:Δsamples:end])
+    close(file)
+
+    return x_data, y_data
+end
+
 function train_don(; n=300, cuda=true, learning_rate=0.001, epochs=400)
     if cuda && has_cuda()
         @info "Training on GPU"
@@ -20,7 +29,7 @@ function train_don(; n=300, cuda=true, learning_rate=0.001, epochs=400)
     opt = ADAM(learning_rate)
 
     m = DeepONet((1024,1024,1024), (1,1024,1024), gelu, gelu) |> device
-    
+
     loss(X, y, sensor) = Flux.Losses.mse(m(X, sensor), y)
     evalcb() = @show(loss(xval, yval, grid))
 

example/Burgers/src/data.jl

@@ -2,6 +2,14 @@ using Burgers
 using Test
 
 @testset "Burgers" begin
-    include("data.jl")
-    include("deeponet.jl")
+    xs, ys = Burgers.get_data(n=1000)
+
+    @test size(xs) == (2, 1024, 1000)
+    @test size(ys) == (1, 1024, 1000)
+
+    learner = Burgers.train(epochs=10)
+    loss = learner.cbstate.metricsepoch[ValidationPhase()][:Loss].values[end]
+    @test loss < 0.1
+
+    # include("deeponet.jl")
 end

example/Burgers/test/data.jl

@@ -2,12 +2,14 @@ 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"
-JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+FluxTraining = "7bf95e4d-ca32-48da-9824-f0dc5310474f"
+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/Burgers/test/runtests.jl

@@ -1,5 +1,5 @@
 ### A Pluto.jl notebook ###
-# v0.16.0
+# v0.19.0
 
 using Markdown
 using InteractiveUtils
@@ -28,10 +28,7 @@ The data is provided by [IBM](https://developer.ibm.com/exchanges/data/all/doubl
 "
 
 # ╔═╡ 5268feee-bda2-4612-9d4c-a1db424a11c7
-data_x, data_y, _, _ = DoublePendulum.preprocess(
-	DoublePendulum.get_data(i=20),
-	ratio=1
-);
+data_x, data_y = DoublePendulum.preprocess(DoublePendulum.get_data(i=20));
 
 # ╔═╡ 4d0b08a4-8a54-41fd-997f-ad54d4c984cd
 m = DoublePendulum.get_model();