Added new subpackage for Sophia.jl

lib/OptimizationSophia/Project.toml

@@ -0,0 +1,28 @@
+name = "OptimizationSophia"
+uuid = "892fee11-dca1-40d6-b698-84ba0d87399a"
+authors = ["paramthakkar123 <[email protected]>"]
+version = "0.1.0"
+
+[deps]
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
+Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
+OptimizationBase = "bca83a33-5cc9-4baa-983d-23429ab6bcbb"
+OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[compat]
+ComponentArrays = "0.15.29"
+Lux = "1.16.0"
+MLUtils = "0.4.8"
+Optimization = "4.5.0"
+OptimizationBase = "2.10.0"
+OrdinaryDiffEqTsit5 = "1.2.0"
+Random = "1.11.0"
+SciMLSensitivity = "7.88.0"
+Test = "1.11.0"
+Zygote = "0.7.10"

src/sophia.jl → lib/OptimizationSophia/src/OptimizationSophia.jl

@@ -1,3 +1,9 @@
+module OptimizationSophia
+
+using OptimizationBase.SciMLBase
+using OptimizationBase: OptimizationCache
+using Optimization
+
 struct Sophia
     η::Float64
     βs::Tuple{Float64, Float64}
@@ -119,3 +125,5 @@ function SciMLBase.__solve(cache::OptimizationCache{
         θ,
         x)
 end
+
+end

lib/OptimizationSophia/test/runtests.jl

@@ -0,0 +1,67 @@
+using OptimizationBase, Optimization
+using OptimizationBase.SciMLBase: solve, OptimizationFunction, OptimizationProblem
+using OptimizationSophia
+using Lux, MLUtils, Random, ComponentArrays
+using SciMLSensitivity
+using Test
+using Zygote
+using OrdinaryDiffEqTsit5
+
+function dudt_(u, p, t)
+    ann(u, p, st)[1] .* u
+end
+
+function newtons_cooling(du, u, p, t)
+    temp = u[1]
+    k, temp_m = p
+    du[1] = dT = -k * (temp - temp_m)
+end
+
+function true_sol(du, u, p, t)
+    true_p = [log(2) / 8.0, 100.0]
+    newtons_cooling(du, u, true_p, t)
+end
+
+function callback(state, l) #callback function to observe training
+    display(l)
+    return l < 1e-2
+end
+
+function predict_adjoint(fullp, time_batch)
+    Array(solve(prob, Tsit5(), p = fullp, saveat = time_batch))
+end
+
+function loss_adjoint(fullp, p)
+    (batch, time_batch) = p
+    pred = predict_adjoint(fullp, time_batch)
+    sum(abs2, batch .- pred)
+end
+
+u0 = Float32[200.0]
+datasize = 30
+tspan = (0.0f0, 1.5f0)
+rng = Random.default_rng()
+
+ann = Lux.Chain(Lux.Dense(1, 8, tanh), Lux.Dense(8, 1, tanh))
+pp, st = Lux.setup(rng, ann)
+pp = ComponentArray(pp)
+
+prob = ODEProblem{false}(dudt_, u0, tspan, pp)
+
+t = range(tspan[1], tspan[2], length = datasize)
+true_prob = ODEProblem(true_sol, u0, tspan)
+ode_data = Array(solve(true_prob, Tsit5(), saveat = t))
+
+k = 10
+train_loader = MLUtils.DataLoader((ode_data, t), batchsize = k)
+
+l1 = loss_adjoint(pp, (train_loader.data[1], train_loader.data[2]))[1]
+
+optfun = OptimizationFunction(loss_adjoint,
+    OptimizationBase.AutoZygote())
+optprob = OptimizationProblem(optfun, pp, train_loader)
+
+res1 = solve(optprob,
+    OptimizationSophia.Sophia(), callback = callback,
+    maxiters = 2000)
+@test 10res1.objective < l1

src/Optimization.jl

@@ -23,7 +23,6 @@ export ObjSense, MaxSense, MinSense
 include("utils.jl")
 include("state.jl")
 include("lbfgsb.jl")
-include("sophia.jl")
 include("auglag.jl")
 
 export solve

test/minibatch.jl

@@ -58,11 +58,6 @@ optfun = OptimizationFunction(loss_adjoint,
     Optimization.AutoZygote())
 optprob = OptimizationProblem(optfun, pp, train_loader)
 
-res1 = Optimization.solve(optprob,
-    Optimization.Sophia(), callback = callback,
-    maxiters = 2000)
-@test 10res1.objective < l1
-
 optfun = OptimizationFunction(loss_adjoint,
     Optimization.AutoForwardDiff())
 optprob = OptimizationProblem(optfun, pp, train_loader)