Merge pull request #9 from EarthyScience/la/lbfgs

lbfgs

Author: lazarusA

examples/Project.toml

@@ -3,12 +3,16 @@ AlgebraOfGraphics = "cbdf2221-f076-402e-a563-3d30da359d67"
 AxisKeys = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
 EasyHybrid = "61bb816a-e6af-4913-ab9e-91bff2e122e3"
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
+OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

examples/Q10_lbfgs.jl

@@ -0,0 +1,68 @@
+using EasyHybrid
+using Lux
+using MLUtils
+using Random
+using Optimization
+using OptimizationOptimisers
+using ComponentArrays
+using GLMakie
+using Random
+using LuxCore
+using CSV, DataFrames
+using Statistics
+using Printf
+
+df = CSV.read("/Users/lalonso/Documents/HybridML/data/Rh_AliceHolt_forcing_filled.csv", DataFrame)
+
+df[!, :Temp] = df[!, :Temp] .- 273.15 # convert to Celsius
+df_forcing = filter(:Respiration_heterotrophic => !isnan, df)
+# df_forcing = df
+ds_k = to_keyedArray(Float32.(df_forcing))
+yobs =  ds_k(:Respiration_heterotrophic)'[:,:]
+
+NN = Lux.Chain(Dense(2, 15, Lux.relu), Dense(15, 15, Lux.relu), Dense(15, 1));
+#? do different initial Q10s
+RbQ10 = RespirationRbQ10(NN, (:Rgpot, :Moist), (:Temp,), 2.5f0) 
+
+data = (ds_k([:Rgpot, :Moist, :Temp]), yobs)
+
+(x_train, y_train), (x_val, y_val) = splitobs(data; at=0.8, shuffle=false)
+dataloader = DataLoader((x_train, y_train), batchsize=512, shuffle=true);
+
+ps, st = LuxCore.setup(Random.default_rng(), RbQ10)
+
+ps_ca = ComponentArray(ps)
+smodel = StatefulLuxLayer{false}(RbQ10, nothing, st)
+# deal with the `Rb` state also here, (; Rb, st), since this is the output from LuxCore.apply.
+# ! note that for now is set to `{false}`.
+
+function callback(state, l)
+    state.iter % 2 == 1 && @printf "Iteration: %5d, Loss: %.6f\n" state.iter l
+    return l < 0.2 ## Terminate if loss is smaller than
+end
+
+function lossfn_optim(ps_ca, data)
+    ds, y = data
+    # ! unpack nan indices here as well
+    ŷ, _ = smodel(ds, ps_ca)
+    return Statistics.mean(abs2, ŷ .- y)
+end
+
+lossfn_optim(ps_ca, (ds_k, yobs))
+
+opt_func = OptimizationFunction(lossfn_optim, Optimization.AutoZygote())
+opt_prob = OptimizationProblem(opt_func, ps_ca, dataloader)
+
+epochs = 25
+res_adam = solve(opt_prob, Optimisers.Adam(0.001); callback, epochs)
+# res_shopia = solve(opt_prob, Optimization.Sophia(); callback, maxiters=epochs)
+
+# ! finetune a bit with L-BFGS
+# ?  LBFGS needs to this `convert(Float64, res_adam.u)` which it fails!
+# ! but there is more, see issue: https://github.com/LuxDL/Lux.jl/issues/1260
+
+# using ForwardDiff
+# opt_func = OptimizationFunction(lossfn_optim, Optimization.AutoForwardDiff())
+# opt_prob2 = remake(opt_prob, u0=res_adam.u)
+opt_prob = OptimizationProblem(opt_func, res_adam.u, dataloader)
+res_lbfgs = solve(opt_prob, Optimization.LBFGS(); callback, maxiters=epochs)