add gradient

Author: foldfelis

example/DoublePendulum/src/DoublePendulum.jl

@@ -15,7 +15,7 @@ function update_model!(model_file_path, model)
     @warn "model updated!"
 end
 
-function train(; loss_bounds=[0.05])
+function train(; loss_bounds=[])
     if has_cuda()
         @info "CUDA is on"
         device = gpu
@@ -25,15 +25,29 @@ function train(; loss_bounds=[0.05])
     end
 
     m = Chain(
-        FourierOperator(1=>64, (6, 64, ), relu),
-        FourierOperator(64=>64, (6, 64, ), relu),
-        FourierOperator(64=>64, (6, 64, ), relu),
-        FourierOperator(64=>1, (6, 64, )),
+        Dense(1, 64, gelu),
+        FourierOperator(64=>64, (12, ), gelu),
+        FourierOperator(64=>64, (12, ), gelu),
+        FourierOperator(64=>64, (12, ), gelu),
+        FourierOperator(64=>64, (12, ), gelu),
+        FourierOperator(64=>64, (12, ), gelu),
+        FourierOperator(64=>64, (12, ), gelu),
+        FourierOperator(64=>64, (12, ), gelu),
+        FourierOperator(64=>64, (12, ), gelu),
+        FourierOperator(64=>64, (12, ), gelu),
+        FourierOperator(64=>64, (12, ), gelu),
+        FourierOperator(64=>64, (12, ), gelu),
+        FourierOperator(64=>64, (12, ), gelu),
+        FourierOperator(64=>64, (12, ), gelu),
+        FourierOperator(64=>64, (12, ), gelu),
+        FourierOperator(64=>64, (12, ), gelu),
+        FourierOperator(64=>64, (12, )),
+        Dense(64, 1)
     ) |> device
 
     loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
 
-    opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-2))
+    opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
 
     loader_train, loader_test = get_dataloader()
 

example/DoublePendulum/src/data.jl

@@ -33,19 +33,18 @@ function get_data(; i=0, n=-1)
     )
     data = (n < 0) ? collect(Matrix(df)') : collect(Matrix(df)')[:, 1:n]
 
-    data /= maximum(data)
-
     return Float32.(data)
 end
 
-function get_dataloader(; i=0, n_train=15734, n_test=2048, Δn=1024, batchsize=100)
-    x = reshape(get_data(; i=i, n=-1), :) # size==(6, 17782)
-    𝐱 = reshape(vcat([x[i:(i+6Δn-1)] for i in 1:6:(length(x)-6(Δn-1))]...), 1, 6, 1024, :)
+function get_dataloader(; i=0, n_train=15733, n_test=2048, Δn=1, batchsize=100)
+    𝐱 = get_data(i=i, n=-1) # size==(6, 17782)
+    ∇𝐱 = 𝐱[:, (1+Δn):end] - 𝐱[:, 1:(end-Δn)]
+    𝐱 = reshape(vcat(𝐱[:, 1:(end-Δn)], ∇𝐱), 1, 12, :)
 
-    𝐱_train, 𝐲_train = 𝐱[:, :, :, 1:(n_train-Δn)], 𝐱[:, :, :, 1+Δn:n_train]
+    𝐱_train, 𝐲_train = 𝐱[:, :, 1:(n_train-1)], 𝐱[:, :, 2:n_train]
     loader_train = Flux.DataLoader((𝐱_train, 𝐲_train), batchsize=batchsize, shuffle=true)
 
-    𝐱_test, 𝐲_test = 𝐱[:, :, :, (end-n_test+1):(end-Δn)], 𝐱[:, :, :, (end-n_test+1+Δn):end]
+    𝐱_test, 𝐲_test = 𝐱[:, :, (end-n_test+1):(end-1)], 𝐱[:, :, (end-n_test+2):end]
     loader_test = Flux.DataLoader((𝐱_test, 𝐲_test), batchsize=batchsize, shuffle=false)
 
     return loader_train, loader_test