docs: add an Optimization.jl tutorial showcasing lazy data movement

docs/make.jl

@@ -23,7 +23,8 @@ pages = [
             "tutorials/beginner/1_Basics.md",
             "tutorials/beginner/2_PolynomialFitting.md",
             "tutorials/beginner/3_SimpleRNN.md",
-            "tutorials/beginner/4_SimpleChains.md"
+            "tutorials/beginner/4_SimpleChains.md",
+            "tutorials/beginner/5_OptimizationIntegration.md"
         ],
         "Intermediate" => [
             "tutorials/intermediate/1_NeuralODE.md",

docs/src/.vitepress/config.mts

@@ -122,7 +122,9 @@ export default defineConfig({
                             { text: 'Julia & Lux for the Uninitiated', link: '/tutorials/beginner/1_Basics' },
                             { text: 'Fitting a Polynomial using MLP', link: '/tutorials/beginner/2_PolynomialFitting' },
                             { text: 'Training a Simple LSTM', link: '/tutorials/beginner/3_SimpleRNN' },
-                            { text: 'MNIST Classification with SimpleChains', link: '/tutorials/beginner/4_SimpleChains' }]
+                            { text: 'MNIST Classification with SimpleChains', link: '/tutorials/beginner/4_SimpleChains' },
+                            { text: 'Fitting with Optimization.jl', link: '/tutorials/beginner/5_OptimizationIntegration' },
+                        ]
                     },
                     {
                         text: 'Intermediate', collapsed: false, items: [

docs/src/tutorials/index.md

@@ -27,11 +27,17 @@ const beginner = [
     caption: "Use SimpleChains.jl as a Backend",
     desc: "Learn how to train small neural networks really fast on CPU."
   },
+  {
+    href: "beginner/5_OptimizationIntegration",
+    src: "../optimization_integration.png",
+    caption: "Fitting with Optimization.jl",
+    desc: "Learn how to use Optimization.jl with Lux (on GPUs)."
+  },
   {
     href: "https://luxdl.github.io/Boltz.jl/stable/tutorials/1_GettingStarted",
     src: "https://production-media.paperswithcode.com/datasets/ImageNet-0000000008-f2e87edd_Y0fT5zg.jpg",
     caption: "Pre-Built Deep Learning Models",
-    desc: "Use Boltz.jl to load pre-built deep learning and scientific machine learning models."
+    desc: "Use Boltz.jl to load pre-built DL and SciML models."
   }
 ];
 

docs/tutorials.jl

@@ -3,7 +3,8 @@ const BEGINNER_TUTORIALS = [
     "Basics/main.jl" => "CUDA",
     "PolynomialFitting/main.jl" => "CUDA",
     "SimpleRNN/main.jl" => "CUDA",
-    "SimpleChains/main.jl" => "CUDA"
+    "SimpleChains/main.jl" => "CUDA",
+    "OptimizationIntegration/main.jl" => "CUDA",
 ]
 const INTERMEDIATE_TUTORIALS = [
     "NeuralODE/main.jl" => "CUDA",

examples/OptimizationIntegration/Project.toml

@@ -0,0 +1,33 @@
+[deps]
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
+Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
+Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
+OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
+OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
+OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
+
+[compat]
+CairoMakie = "0.12.10"
+ComponentArrays = "0.15.17"
+InteractiveUtils = "<0.0.1, 1"
+IterTools = "1.10"
+Literate = "2.19"
+Lux = "1"
+LuxCUDA = "0.3.3"
+MLUtils = "0.4.4"
+Optimization = "3.28.0"
+OptimizationOptimJL = "0.3.2"
+OptimizationOptimisers = "0.2.1"
+OrdinaryDiffEqTsit5 = "1.1.0"
+Printf = "1.10"
+Random = "1.10"
+SciMLSensitivity = "7.67.0"

examples/OptimizationIntegration/main.jl

@@ -0,0 +1,158 @@
+# # Training Lux Models using Optimization.jl
+
+# Lux's native [Training.TrainState](@ref) is a great API for gradient-based learning of
+# neural networks, however, it is geared towards using `Optimisers.jl` as the backend.
+# However, often times we want to train the neural networks with other optimization methods
+# like BFGS, LBFGS, etc. In this tutorial, we will show how to train Lux models with
+# Optimization.jl that provides a simple unified interface to various optimization methods.
+
+# We will base our tutorial on the minibatching tutorial from the official
+# [Optimization.jl](https://docs.sciml.ai/Optimization/stable/tutorials/minibatch/) docs.
+
+# !!! note "Neural ODE"
+#
+#     This tutorial uses a Neural ODE, however, we won't discuss that part in this tutorial.
+#     Please refer to the Neural ODE tutorial for more information.
+
+# ## Imports packages
+
+using Lux, Optimization, OptimizationOptimisers, OptimizationOptimJL, OrdinaryDiffEqTsit5,
+      SciMLSensitivity, Random, MLUtils, IterTools, CairoMakie, ComponentArrays, Printf
+using LuxCUDA
+
+const gdev = gpu_device()
+const cdev = cpu_device()
+
+# ## Generate some training data
+
+function lotka_volterra(du, u, p, t)
+    x, y = u
+    α, β, δ, γ = p
+    du[1] = α * x - β * x * y
+    du[2] = -δ * y + γ * x * y
+    return nothing
+end
+
+u0 = [1.0f0, 1.0f0]
+
+datasize = 32
+tspan = (0.0f0, 2.0f0)
+
+const t = range(tspan[1], tspan[2]; length=datasize)
+true_prob = ODEProblem(lotka_volterra, u0, (tspan[1], tspan[2]), [1.5, 1.0, 3.0, 1.0])
+const ode_data = Array(solve(true_prob, Tsit5(); saveat=t))
+
+begin
+    fig = Figure()
+    ax = CairoMakie.Axis(fig[1, 1])
+    lines!(ax, t, ode_data[1, :]; label=L"u_1(t)", color=:blue, linestyle=:dot, linewidth=4)
+    lines!(ax, t, ode_data[2, :]; label=L"u_2(t)", color=:red, linestyle=:dot, linewidth=4)
+    axislegend(ax; position=:lt)
+    fig
+end
+
+# ## Define the DataLoader
+
+# We will define the DataLoader to batch over the data, additionally we will pipe it through
+# the `gdev` device to move the data to the GPU on each iteration.
+
+# By default `gdev` will move all objects to the GPU. But we don't want to move the time
+# vector to the GPU. So we will wrap it in a struct.
+struct TimeWrapper{T}
+    t::T
+end
+
+Base.length(t::TimeWrapper) = length(t.t)
+
+Base.getindex(t::TimeWrapper, i) = TimeWrapper(t.t[i])
+
+dataloader = DataLoader((ode_data, TimeWrapper(t)); batchsize=8) |> gdev
+
+# ## Training the model
+
+# Here we are using different optimization methods for demonstration purposes. This problem
+# is trivial enough to not require this.
+
+# Optimization.jl requires an abstract array as the parameters, hence we will construct a
+# `ComponentArray` to store the parameters.
+
+# !!! note "Parameter Estimation vs State Estimation"
+#
+#     Optimization.jl performs state estimation, which effectively means for a function
+#     `f(u, p)`, it is trying to compute the optimal `u` for a given `p`. This terminology
+#     might be confusing to ML practitioners, since in the ML world, we usually do parameter
+#     estimation. This effectively means that the `u` in Optimization.jl corresponds to our
+#     model parameters that is being optimized.
+
+function train_model(dataloader)
+    model = Chain(Dense(2, 32, tanh), Dense(32, 32, tanh), Dense(32, 2))
+    ps, st = Lux.setup(Random.default_rng(), model)
+
+    ps_ca = ComponentArray(ps) |> gdev
+    st = st |> gdev
+
+    function callback(state, l)
+        state.iter % 25 == 1 && @printf "Iteration: %5d, Loss: %.6e\n" state.iter l
+        return l < 1e-8 ## Terminate if loss is small
+    end
+
+    smodel = StatefulLuxLayer{true}(model, nothing, st)
+
+    function loss_adjoint(θ, u_batch, t_batch)
+        t_batch = t_batch.t
+        u0 = u_batch[:, 1]
+        dudt(u, p, t) = smodel(u, p)
+        prob = ODEProblem(dudt, u0, (t_batch[1], t_batch[end]), θ)
+        pred = convert(AbstractArray, solve(prob, Tsit5(); saveat=t_batch))
+        return MSELoss()(pred, u_batch)
+    end
+
+    ## Define the Optimization Function that takes in the optimization state (our parameters)
+    ## and optimization parameters (nothing in our case) and data from the dataloader and
+    ## returns the loss.
+    opt_func = OptimizationFunction(
+        (θ, _, u_batch, t_batch) -> loss_adjoint(θ, u_batch, t_batch),
+        Optimization.AutoZygote())
+    opt_prob = OptimizationProblem(opt_func, ps_ca)
+
+    nepcohs = 25
+    res_adam = solve(
+        opt_prob, Optimisers.Adam(0.001), ncycle(dataloader, nepcohs); callback)
+
+    ## Let's finetune a bit with L-BFGS
+    opt_prob = remake(opt_prob; u0=res_adam.u)
+    res_lbfgs = solve(opt_prob, LBFGS(), ncycle(dataloader, nepcohs); callback)
+
+    ## Now that we have a good fit, let's train it on the entire dataset without
+    ## Minibatching. We need to do this since ODE solves can lead to accumulated errors if
+    ## the model was trained on individual parts (without a data-shooting approach).
+    opt_func = OptimizationFunction(
+        (θ, _) -> loss_adjoint(θ, gdev(ode_data), TimeWrapper(t)),
+        Optimization.AutoZygote())
+    opt_prob = OptimizationProblem(opt_func, res_lbfgs.u)
+
+    res = solve(opt_prob, Optimisers.Adam(0.005); maxiters=500, callback)
+
+    return StatefulLuxLayer{true}(model, res.u, smodel.st)
+end
+
+trained_model = train_model(dataloader)
+nothing #hide
+
+# ## Plotting the results
+
+dudt(u, p, t) = trained_model(u, p)
+prob = ODEProblem(dudt, gdev(u0), (tspan[1], tspan[2]), trained_model.ps)
+sol = solve(prob, Tsit5(); saveat=t)
+pred = convert(AbstractArray, sol) |> cdev
+
+begin
+    fig = Figure()
+    ax = CairoMakie.Axis(fig[1, 1])
+    lines!(ax, t, ode_data[1, :]; label=L"u_1(t)", color=:blue, linestyle=:dot, linewidth=4)
+    lines!(ax, t, ode_data[2, :]; label=L"u_2(t)", color=:red, linestyle=:dot, linewidth=4)
+    lines!(ax, t, pred[1, :]; label=L"\hat{u}_1(t)", color=:blue, linewidth=4)
+    lines!(ax, t, pred[2, :]; label=L"\hat{u}_2(t)", color=:red, linewidth=4)
+    axislegend(ax; position=:lt)
+    fig
+end