Merge pull request #74 from SciML/gpu

Add downstream tests and fix GPU

Author: ChrisRackauckas

.travis.yml

@@ -10,6 +10,9 @@ os:
 julia:
   - 1
 #  - nightly
+env:
+  - GROUP=Core
+  - GROUP=Downstream
 notifications:
   email: false
 jobs:

Project.toml

@@ -48,11 +48,12 @@ CMAEvolutionStrategy = "8d3b24bd-414e-49e0-94fb-163cc3a3e411"
 DiffEqFlux = "aae7a2af-3d4f-5e19-a356-7da93b79d9d0"
 Evolutionary = "86b6b26d-c046-49b6-aa0b-5f0f74682bd6"
 IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 NLopt = "76087f3c-5699-56af-9a33-bf431cd00edd"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["BlackBoxOptim", "Evolutionary", "DiffEqFlux", "IterTools", "OrdinaryDiffEq", "NLopt", "CMAEvolutionStrategy", "Plots" ,"SafeTestsets", "Test"]
+test = ["BlackBoxOptim", "Evolutionary", "DiffEqFlux", "IterTools", "OrdinaryDiffEq", "NLopt", "CMAEvolutionStrategy", "Plots", "Pkg", "SafeTestsets", "Test"]

src/function.jl

@@ -37,7 +37,7 @@ function instantiate_function(f, x, ::AbstractADType, p, num_cons = 0)
                          cons_j,cons_h)
 end
 
-function instantiate_function(f, x, ::AutoForwardDiff{_chunksize}, p, num_cons = 0) where _chunksize
+function instantiate_function(f::OptimizationFunction{true}, x, ::AutoForwardDiff{_chunksize}, p, num_cons = 0) where _chunksize
 
     chunksize = _chunksize === nothing ? default_chunk_size(length(x)) : _chunksize
 

src/solve.jl

@@ -12,6 +12,7 @@ function DiffEqBase.solve(prob::OptimizationProblem, opt, args...;kwargs...)
 	__solve(prob, opt, args...; kwargs...)
 end
 
+#=
 function update!(x::AbstractArray, x̄::AbstractArray{<:ForwardDiff.Dual})
   x .-= x̄
 end
@@ -31,6 +32,7 @@ end
 function update!(opt, xs::Flux.Zygote.Params, gs)
 	update!(opt, xs[1], gs)
 end
+=#
 
 maybe_with_logger(f, logger) = logger === nothing ? f() : Logging.with_logger(f, logger)
 
@@ -62,7 +64,10 @@ macro withprogress(progress, exprs...)
 	end |> esc
   end
 
-function __solve(prob::OptimizationProblem, opt, data = DEFAULT_DATA;cb = (args...) -> (false), maxiters::Number = 1000, progress = true, save_best = true, kwargs...)
+function __solve(prob::OptimizationProblem, opt, data = DEFAULT_DATA;
+				 cb = (args...) -> (false), maxiters::Number = 1000,
+				 progress = true, save_best = true, kwargs...)
+
 	if maxiters <= 0.0
 		error("The number of maxiters has to be a non-negative and non-zero number.")
 	else
@@ -76,7 +81,7 @@ function __solve(prob::OptimizationProblem, opt, data = DEFAULT_DATA;cb = (args.
 
 	if data != DEFAULT_DATA
 		maxiters = length(data)
-	else 
+	else
 		data = take(data, maxiters)
 	end
 
@@ -90,8 +95,10 @@ function __solve(prob::OptimizationProblem, opt, data = DEFAULT_DATA;cb = (args.
 
 	@withprogress progress name="Training" begin
 	  for (i,d) in enumerate(data)
-		gs = prob.f.adtype isa AutoFiniteDiff ? Array{Number}(undef,length(θ)) : DiffResults.GradientResult(θ)
-		f.grad(gs, θ, d...) 
+	    gs = Flux.Zygote.gradient(ps) do
+  		  x = prob.f(θ,prob.p, d...)
+  		  first(x)
+  		end
 		x = f.f(θ, prob.p, d...)
 		cb_call = cb(θ, x...)
 		if !(typeof(cb_call) <: Bool)
@@ -101,7 +108,7 @@ function __solve(prob::OptimizationProblem, opt, data = DEFAULT_DATA;cb = (args.
 		end
 		msg = @sprintf("loss: %.3g", x[1])
 		progress && ProgressLogging.@logprogress msg i/maxiters
-		update!(opt, ps, prob.f.adtype isa AutoFiniteDiff ? gs : DiffResults.gradient(gs))
+		Flux.update!(opt, ps, gs)
 
 		if save_best
 		  if first(x) < first(min_err)  #found a better solution
@@ -215,7 +222,7 @@ function __solve(prob::OptimizationProblem, opt::Union{Optim.Fminbox,Optim.SAMIN
 	  	if !(typeof(cb_call) <: Bool)
 			error("The callback should return a boolean `halt` for whether to stop the optimization process.")
 		end
-		cur, state = iterate(data, state)  
+		cur, state = iterate(data, state)
 	  	cb_call
 	end
 

test/downstream/Project.toml

@@ -0,0 +1,4 @@
+[deps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+DiffEqFlux = "aae7a2af-3d4f-5e19-a356-7da93b79d9d0"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"

test/downstream/gpu_neural_ode.jl

@@ -0,0 +1,54 @@
+using DiffEqFlux, OrdinaryDiffEq, Flux, CUDA
+CUDA.allowscalar(false) # Makes sure no slow operations are occuring
+
+# Generate Data
+u0 = Float32[2.0; 0.0]
+datasize = 30
+tspan = (0.0f0, 1.5f0)
+tsteps = range(tspan[1], tspan[2], length = datasize)
+function trueODEfunc(du, u, p, t)
+    true_A = [-0.1 2.0; -2.0 -0.1]
+    du .= ((u.^3)'true_A)'
+end
+prob_trueode = ODEProblem(trueODEfunc, u0, tspan)
+# Make the data into a GPU-based array if the user has a GPU
+ode_data = gpu(solve(prob_trueode, Tsit5(), saveat = tsteps))
+
+
+dudt2 = FastChain((x, p) -> x.^3,
+                  FastDense(2, 50, tanh),
+                  FastDense(50, 2))
+u0 = Float32[2.0; 0.0] |> gpu
+p = initial_params(dudt2) |> gpu
+prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(), saveat = tsteps)
+
+function predict_neuralode(p)
+  gpu(prob_neuralode(u0,p))
+end
+function loss_neuralode(p)
+    pred = predict_neuralode(p)
+    loss = sum(abs2, ode_data .- pred)
+    return loss, pred
+end
+# Callback function to observe training
+list_plots = []
+iter = 0
+callback = function (p, l, pred; doplot = false)
+  global list_plots, iter
+  if iter == 0
+    list_plots = []
+  end
+  iter += 1
+  display(l)
+  # plot current prediction against data
+  plt = scatter(tsteps, Array(ode_data[1,:]), label = "data")
+  scatter!(plt, tsteps, Array(pred[1,:]), label = "prediction")
+  push!(list_plots, plt)
+  if doplot
+    display(plot(plt))
+  end
+  return false
+end
+result_neuralode = DiffEqFlux.sciml_train(loss_neuralode, p,
+                                          ADAM(0.05), cb = callback,
+                                          maxiters = 300)

test/runtests.jl

@@ -1,10 +1,30 @@
-using SafeTestsets
-
-println("Rosenbrock Tests")
-@safetestset "Rosenbrock" begin include("rosenbrock.jl") end
-println("AD Tests")
-@safetestset "AD Tests" begin include("ADtests.jl") end
-println("Mini batching Tests")
-@safetestset "Mini batching" begin include("minibatch.jl") end
-println("DiffEqFlux Tests")
-@safetestset "DiffEqFlux" begin include("diffeqfluxtests.jl") end
+using SafeTestsets, Pkg
+
+const GROUP = get(ENV, "GROUP", "All")
+const is_APPVEYOR = Sys.iswindows() && haskey(ENV,"APPVEYOR")
+const is_TRAVIS = haskey(ENV,"TRAVIS")
+
+function activate_downstream_env()
+    Pkg.activate("downstream")
+    Pkg.develop(PackageSpec(path=dirname(@__DIR__)))
+    Pkg.instantiate()
+end
+
+@time begin
+if GROUP == "All" || GROUP == "Core"
+    @safetestset "Rosenbrock" begin include("rosenbrock.jl") end
+    @safetestset "AD Tests" begin include("ADtests.jl") end
+    @safetestset "Mini batching" begin include("minibatch.jl") end
+    @safetestset "DiffEqFlux" begin include("diffeqfluxtests.jl") end
+end
+
+if !is_APPVEYOR && GROUP == "Downstream"
+    activate_downstream_env()
+    Pkg.test("DiffEqFlux")
+end
+
+if !is_APPVEYOR && GROUP == "GPU"
+    activate_downstream_env()
+    @safetestset "DiffEqFlux GPU" begin include("downstream/gpu_neural_ode.jl") end
+end
+end