Update Neural ODE benchmark

Author: utkarsh530

benchmarks/adam_opt.jl

@@ -1,6 +1,8 @@
 using SimpleChains,
     StaticArrays, OrdinaryDiffEq, SciMLSensitivity, Optimization, OptimizationFlux, Plots
 
+#device!(2)
+# Get Tesla V100S
 u0 = @SArray Float32[2.0, 0.0]
 datasize = 30
 tspan = (0.0f0, 1.5f0)
@@ -23,10 +25,10 @@ p_nn = SimpleChains.init_params(sc)
 
 f(u, p, t) = sc(u, p)
 
-prob_nn = ODEProblem(f, u0, tspan)
+sprob_nn = ODEProblem(f, u0, tspan)
 
 function predict_neuralode(p)
-    Array(solve(prob_nn,
+    Array(solve(sprob_nn,
         Tsit5();
         p = p,
         saveat = tsteps,
@@ -56,6 +58,8 @@ optprob = Optimization.OptimizationProblem(optf, p_nn)
 @time res = Optimization.solve(optprob, ADAM(0.05), maxiters = 100)
 @show res.objective
 
+@benchmark Optimization.solve(optprob, ADAM(0.05), maxiters = 100)
+
 ## PSOGPU stuff
 
 function nn_fn(u::T, p, t)::T where {T}
@@ -78,12 +82,15 @@ function loss(u, p)
     sum(abs2, data .- pred)
 end
 
-lb = SVector{length(p_static), eltype(p_static)}(fill(eltype(p_static)(-10),
-    length(p_static))...)
-ub = SVector{length(p_static), eltype(p_static)}(fill(eltype(p_static)(10),
-    length(p_static))...)
+# lb = SVector{length(p_static), eltype(p_static)}(fill(eltype(p_static)(-10),
+#     length(p_static))...)
+# ub = SVector{length(p_static), eltype(p_static)}(fill(eltype(p_static)(10),
+#     length(p_static))...)
+
+lb = @SArray fill(Float32(-Inf), length(p_static))
+ub = @SArray fill(Float32(Inf), length(p_static))
 
-optprob = OptimizationProblem(loss, prob_nn.p[2], (prob_nn, tsteps); lb = lb, ub = ub)
+soptprob = OptimizationProblem(loss, prob_nn.p[2], (prob_nn, tsteps); lb = lb, ub = ub)
 
 using PSOGPU
 using CUDA
@@ -93,27 +100,58 @@ using Adapt
 backend = CUDABackend()
 
 ## Initialize Particles
-gbest, particles = PSOGPU.init_particles(optprob,
+gbest, particles = PSOGPU.init_particles(soptprob,
     ParallelSyncPSOKernel(n_particles; backend = CUDABackend()),
     typeof(prob.u0))
 
 gpu_data = adapt(backend,
     [SVector{length(prob_nn.u0), eltype(prob_nn.u0)}(@view data[:, i])
      for i in 1:length(tsteps)])
 
-gpu_particles = adapt(backend, particles)
-
 CUDA.allowscalar(false)
 
 function prob_func(prob, gpu_particle)
     return remake(prob, p = (prob.p[1], gpu_particle.position))
 end
 
+gpu_particles = adapt(backend, particles)
+
+losses = adapt(backend, ones(eltype(prob.u0), (1, n_particles)))
+
+solver_cache = (; losses, gpu_particles, gpu_data, gbest)
+
 @time gsol = PSOGPU.parameter_estim_ode!(prob_nn,
-    gpu_particles,
-    gbest,
-    gpu_data,
+    solver_cache,
     lb,
-    ub; saveat = tsteps, dt = 0.1f0, prob_func = prob_func, maxiters = 100)
+    ub;
+    saveat = tsteps,
+    dt = 0.1f0,
+    prob_func = prob_func,
+    maxiters = 100)
+
+@benchmark PSOGPU.parameter_estim_ode!($prob_nn,
+    $(deepcopy(solver_cache)),
+    $lb,
+    $ub;
+    saveat = tsteps,
+    dt = 0.1f0,
+    prob_func = prob_func,
+    maxiters = 100)
+
+@show gsol.best
+
+using Plots
+
+function predict_neuralode(p)
+    Array(solve(prob_nn, Tsit5(); p = p, saveat = tsteps))
+end
+
+using Plots
+
+plt = scatter(tsteps, data[1, :], label = "data")
+
+pred_pso = predict_neuralode((sc, gsol.position))
+scatter!(plt, tsteps, pred[1, :], label = "PSO prediction")
 
-@show gsol.cost
+pred_adam = predict_neuralode((sc, res.u))
+scatter!(plt, tsteps, pred_adam[1, :], label = "Adam prediction")

src/ode_pso.jl

@@ -45,23 +45,17 @@ function default_prob_func(prob, gpu_particle)
     return remake(prob, p = gpu_particle.position)
 end
 
-function parameter_estim_ode!(prob::ODEProblem,
-        gpu_particles,
-        gbest,
-        data,
+function parameter_estim_ode!(prob::ODEProblem, cache,
         lb,
         ub;
         ode_alg = GPUTsit5(),
         prob_func = default_prob_func,
         w = 0.72980f0,
         wdamp = 1.0f0,
         maxiters = 100, kwargs...)
+    (losses, gpu_particles, gpu_data, gbest) = cache
     backend = get_backend(gpu_particles)
     update_states! = PSOGPU._update_particle_states!(backend)
-
-    losses = KernelAbstractions.allocate(backend,
-        eltype(prob.u0),
-        (1, length(gpu_particles)))
     update_costs! = PSOGPU._update_particle_costs!(backend)
 
     improb = make_prob_compatible(prob)
@@ -80,16 +74,11 @@ function parameter_estim_ode!(prob::ODEProblem,
             prob,
             ode_alg; kwargs...)
 
-        sum!(losses, (map(x -> sum(x .^ 2), data .- us)))
+        sum!(losses, (map(x -> sum(x .^ 2), gpu_data .- us)))
 
         update_costs!(losses, gpu_particles; ndrange = length(losses))
 
-        best_particle = minimum(gpu_particles,
-            init = PSOGPU.SPSOParticle(gbest.position,
-                gbest.position,
-                gbest.cost,
-                gbest.position,
-                gbest.cost))
+        best_particle = minimum(gpu_particles)
 
         gbest = PSOGPU.SPSOGBest(best_particle.best_position, best_particle.best_cost)
         w = w * wdamp