Для обучения модели на GPU я использую
dudt = Chain(Dense(3,100,tanh),
Dense(100,3)) |> gpu
против
Обучение ЦП
dudt = FastChain(
FastDense(3,100,tanh),
FastDense(100,3))
За 1000 итераций Fastchain на несколько порядков быстрее, чем работает графический процессор Tesla K40 c. Это ожидаемое поведение? В противном случае, я мог бы сделать что-то не так с реализацией модели на графических процессорах. MWE для реализации GPU выглядит следующим образом:
function lorenz(du,u,p,t)
σ = p[1]; ρ = p[2]; β = p[3]
du[1] = σ*(u[2]-u[1])
du[2] = u[1]*(ρ-u[3]) - u[2]
du[3] = u[1]*u[2] - β*u[3]
return
end
u0 = Float32[1.0,0.0,0.0]
tspan = (0.0,1.0)
para = [10.0,28.0,8/3]
prob = ODEProblem(lorenz, u0, tspan, para)
t = range(tspan[1],tspan[2],length=101)
ode_data = Array(solve(prob,Tsit5(),saveat=t))
ode_data = cu(ode_data)
u0train = [1.0,0.0,0.0] |> gpu
tspantrain = (0.0,1.0)
ttrain = range(tspantrain[1],tspantrain[2],length=101)
dudt = Chain(Dense(3,100,tanh),
Dense(100,3)) |> gpu
n_ode = NeuralODE((dudt),tspantrain,Tsit5(),saveat=ttrain)
function predict_n_ode(p)
n_ode(u0train,p)
end
function loss_n_ode(p)
pred = predict_n_ode(p) |> gpu
loss = sum(abs2, pred .- ode_data)
loss,pred
end
res1 = DiffEqFlux.sciml_train(loss_n_ode, n_ode.p, ADAM(0.01), cb=cb, maxiters = 1000)