Как я могу получить доступ к обученным параметрам нейронного ODE в Julia?

Question

Я пытаюсь подогнать одно Neural ODE к временному ряду, используя DiffEqFlux Джулии. Вот мой код:

u0 = Float32[2.;0]
train_size = 15
tspan_train = (0.0f0,0.75f0)

function trueODEfunc(du,u,p,t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u.^3)'true_A)'
end

t_train = range(tspan_train[1],tspan_train[2],length = train_size)
prob = ODEProblem(trueODEfunc, u0, tspan_train)
ode_data_train = Array(solve(prob, Tsit5(),saveat=t_train))

dudt = Chain(
            Dense(2,50,tanh),
            Dense(50,2))
ps = Flux.params(dudt)
n_ode = NeuralODE(dudt, tspan_train, Tsit5(), saveat = t_train, reltol=1e-7, abstol=1e-9)

**n_ode.p**

function predict_n_ode(p)
    n_ode(u0,p)
end
function loss_n_ode(p)
    pred = predict_n_ode(p)
    loss = sum(abs2, ode_data_train .- pred)
    loss,pred
end

final_p = []
losses = []
cb = function(p,l,pred)
    display(l)
    display(p)
    push!(final_p, p)
    push!(losses,l)
    pl = scatter(t_train, ode_data_train[1,:],label="data")
    scatter!(pl,t_train,pred[1,:],label="prediction")
    display(plot(pl))
end

DiffEqFlux.sciml_train!(loss_n_ode, n_ode.p, ADAM(0.05), cb = cb, maxiters = 100)

**n_ode.p**

Проблема в том, что звонок n_ode.p (или Flux.params(dudt)) до и после функции поезда возвращают мне сохраненные значения. Я ожидал получить последние обновленные значения от обучения. Вот почему я создал массив, чтобы собрать все значения параметров во время обучения, а затем получить к нему доступ для получения обновленных параметров.

Я что-то не так делаю в коде? Функция поезда обновляет параметры автоматически? Если нет, то как это сделать?

Заранее спасибо!

Answer 1

В результате получается объект с наилучшими параметрами. Вот полный пример:

using DiffEqFlux, OrdinaryDiffEq, Flux, Optim, Plots

u0 = Float32[2.; 0.]
datasize = 30
tspan = (0.0f0,1.5f0)

function trueODEfunc(du,u,p,t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u.^3)'true_A)'
end
t = range(tspan[1],tspan[2],length=datasize)
prob = ODEProblem(trueODEfunc,u0,tspan)
ode_data = Array(solve(prob,Tsit5(),saveat=t))

dudt2 = FastChain((x,p) -> x.^3,
            FastDense(2,50,tanh),
            FastDense(50,2))
n_ode = NeuralODE(dudt2,tspan,Tsit5(),saveat=t)

function predict_n_ode(p)
  n_ode(u0,p)
end

function loss_n_ode(p)
    pred = predict_n_ode(p)
    loss = sum(abs2,ode_data .- pred)
    loss,pred
end

loss_n_ode(n_ode.p) # n_ode.p stores the initial parameters of the neural ODE

cb = function (p,l,pred;doplot=false) #callback function to observe training
  display(l)
  # plot current prediction against data
  if doplot
    pl = scatter(t,ode_data[1,:],label="data")
    scatter!(pl,t,pred[1,:],label="prediction")
    display(plot(pl))
  end
  return false
end

# Display the ODE with the initial parameter values.
cb(n_ode.p,loss_n_ode(n_ode.p)...)

res1 = DiffEqFlux.sciml_train(loss_n_ode, n_ode.p, ADAM(0.05), cb = cb, maxiters = 300)
cb(res1.minimizer,loss_n_ode(res1.minimizer)...;doplot=true)
res2 = DiffEqFlux.sciml_train(loss_n_ode, res1.minimizer, LBFGS(), cb = cb)
cb(res2.minimizer,loss_n_ode(res2.minimizer)...;doplot=true)

# result is res2 as an Optim.jl object
# res2.minimizer are the best parameters
# res2.minimum is the best loss

В конце концов, sciml_train функция возвращает объект результата, который содержит информацию об оптимизации, включая окончательные параметры как .minimizer.