Introduce documentation on how to use mini-batches with Lux
Nobody has claimed this yet.
- Dominant language
- Julia
- Stars
- 839
- Forks
- 100
- Avg merge
- 20h 43m
- Merged PRs (30d)
- 38
Description
Following the documentation in https://docs.sciml.ai/Optimization/stable/tutorials/minibatch/#Data-Iterators-and-Minibatching.
I tried to replace the Flux library with Lux as:
using Lux, Optimization, OptimizationOptimisers, OrdinaryDiffEq, SciMLSensitivity
using StableRNGs
import MLUtils: DataLoader
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
rng = StableRNG(1111)
ann = Lux.Chain(Dense(1, 8, tanh), Dense(8, 1, tanh))
pp, st = Lux.setup(rng, ann)
function dudt_(u, p, t)
ann(u,p,st)[1] .* u
end
callback = function (p, l) #callback function to observe training
display(l)
return false
end
u0 = Float32[200.0]
datasize = 30
tspan = (0.0f0, 1.5f0)
t = range(tspan[1], tspan[2], length = datasize)
true_prob = ODEProblem(true_sol, u0, tspan)
ode_data = Array(solve(true_prob, Tsit5(), saveat = t))
prob = ODEProblem{false}(dudt_, u0, tspan, pp)
function predict_adjoint(fullp, time_batch)
Array(solve(prob, Tsit5(), p = fullp, saveat = time_batch))
end
function loss_adjoint(fullp, batch, time_batch)
pred = predict_adjoint(fullp, time_batch)
sum(abs2, batch .- pred)
end
k = 10
# Pass the data for the batches as separate vectors wrapped in a tuple
train_loader = DataLoader((ode_data, t), batchsize = k)
numEpochs = 300
l1 = loss_adjoint(pp, train_loader.data[1], train_loader.data[2])[1]
optfun = OptimizationFunction((θ, p, batch, time_batch) -> loss_adjoint(θ, batch,
time_batch),
Optimization.AutoZygote())
optprob = OptimizationProblem(optfun, pp)
using IterTools: ncycle
res1 = Optimization.solve(optprob, Optimisers.ADAM(0.05), ncycle(train_loader, numEpochs),
callback = callback)
However, this yields an error:
ERROR: MethodError: no method matching copy(::NamedTuple{(:layer_1, :layer_2), Tuple{NamedTuple{(:weight, :bias), Tuple{Matrix{Float32}, Matrix{Float32}}}, NamedTuple{(:weight, :bias), Tuple{Matrix{Float32}, Matrix{Float32}}}}})
Closest candidates are:
copy(::Union{DiffEqNoiseProcess.BoxWedgeTail, DiffEqNoiseProcess.NoiseApproximation, DiffEqNoiseProcess.NoiseGrid, DiffEqNoiseProcess.NoiseWrapper, DiffEqNoiseProcess.VirtualBrownianTree})
@ DiffEqNoiseProcess ~/.julia/packages/DiffEqNoiseProcess/VQe6Y/src/copy_noise_types.jl:55
copy(::Random123.Threefry4x{T, R}) where {T, R}
@ Random123 ~/.julia/packages/Random123/u5oEp/src/threefry.jl:266
copy(::Zygote.Buffer)
@ Zygote ~/.julia/packages/Zygote/JeHtr/src/tools/buffer.jl:64
Would you consider introducing documentation on how to use mini-batches with Lux? This is a library used for universal differential equations and it would be useful to use this approach to train UODEs with different initial conditions/ parameters.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the linked Optimization mini-batching tutorial and the issue's Lux, MLUtils.DataLoader, OptimizationFunction, and Optimization.solve example. Reproduce the reported NamedTuple copy error, then document a working mini-batch workflow for Lux-based universal differential equations, including the relevant data and parameter handling. Done means a clear example lets users train across batches without the reported failure.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- documentation, machine-learning
- Issue type
- Documentation
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100