Итераторы данных и мини-пакеты

Страница в процессе перевода.

Задачу оптимизации можно решать пакетами, используя объект MLUtils.DataLoader, который передается в Optimization.solve с ncycles. Все данные для пакетов необходимо передавать в виде кортежа векторов.

Note В этом примере применяется пакет OptimizationOptimisers.jl. Подробные сведения об установке и использовании см. на странице Optimisers.jl.

using Lux, Optimization, OptimizationOptimisers, OrdinaryDiffEq, SciMLSensitivity, MLUtils,
      Random, ComponentArrays

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

model = Chain(Dense(1, 32, tanh), Dense(32, 1))
ps, st = Lux.setup(Random.default_rng(), model)
ps_ca = ComponentArray(ps)
smodel = StatefulLuxLayer{true}(model, nothing, st)

function dudt_(u, p, t)
    smodel(u, p) .* u
end

function callback(state, l) #функция обратного вызова для наблюдения за обучением
    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, ps_ca)

function predict_adjoint(fullp, time_batch)
    Array(solve(prob, Tsit5(), p = fullp, saveat = time_batch))
end

function loss_adjoint(fullp, data)
    batch, time_batch = data
    pred = predict_adjoint(fullp, time_batch)
    sum(abs2, batch .- pred)
end

k = 10
# Передаем данные для пакетов в виде отдельных векторов, заключенных в кортеж
train_loader = MLUtils.DataLoader((ode_data, t), batchsize = k)

numEpochs = 300
l1 = loss_adjoint(ps_ca, train_loader.data)[1]

optfun = OptimizationFunction(
    loss_adjoint,
    Optimization.AutoZygote())
optprob = OptimizationProblem(optfun, ps_ca, train_loader)
using IterTools: ncycle
res1 = Optimization.solve(
    optprob, Optimisers.ADAM(0.05); callback = callback, epochs = 1000)

retcode: Default
u: ComponentVector{Float32}(layer_1 = (weight = Float32[-0.5693787; -0.37341732; … ; 0.65470386; 1.1405532;;], bias = Float32[-0.5159028, 0.73251, 2.210856, -1.5252011, 3.3779244, -0.6406089, 3.5126567, -2.1221583, -2.4785213, 0.9479559  …  -1.4701205, 1.4514927, 3.746335, -1.6434731, 1.4498546, -1.185644, -2.9537063, 1.2789421, 0.826725, 0.0922288]), layer_2 = (weight = Float32[0.076110855 0.068753846 … 0.22415866 0.028837912], bias = Float32[-0.43753505]))