Merge pull request #71 from SciML/dataarg

Clean up data arg checking

Author: Vaibhavdixit02

Project.toml

@@ -45,10 +45,14 @@ julia = "1.3"
 [extras]
 BlackBoxOptim = "a134a8b2-14d6-55f6-9291-3336d3ab0209"
 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"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["BlackBoxOptim", "Evolutionary", "NLopt", "CMAEvolutionStrategy", "SafeTestsets", "Test"]
+test = ["BlackBoxOptim", "Evolutionary", "DiffEqFlux", "IterTools", "OrdinaryDiffEq", "NLopt", "CMAEvolutionStrategy", "Plots" ,"SafeTestsets", "Test"]

src/solve.jl

@@ -62,7 +62,7 @@ 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
@@ -74,14 +74,10 @@ function __solve(prob::OptimizationProblem, opt, _data = DEFAULT_DATA;cb = (args
 	θ = copy(prob.u0)
 	ps = Flux.params(θ)
 
-	if _data == DEFAULT_DATA && maxiters == typemax(Int)
-		error("For Flux optimizers, either a data iterator must be provided or the `maxiters` keyword argument must be set.")
-	  elseif _data == DEFAULT_DATA && maxiters != typemax(Int)
-		data = Iterators.repeated((), maxiters)
-	  elseif maxiters != typemax(Int)
-		data = take(_data, maxiters)
-	  else
-		data = _data
+	if data != DEFAULT_DATA
+		maxiters = length(data)
+	else 
+		data = take(data, maxiters)
 	end
 
 	t0 = time()
@@ -158,6 +154,11 @@ decompose_trace(trace) = trace
 
 function __solve(prob::OptimizationProblem, opt::Optim.AbstractOptimizer, data = DEFAULT_DATA;cb = (args...) -> (false), maxiters::Number = 1000, kwargs...)
   	local x, cur, state
+
+	if data != DEFAULT_DATA
+		maxiters = length(data)
+	end
+
 	cur, state = iterate(data)
 
 	function _cb(trace)
@@ -202,6 +203,11 @@ end
 
 function __solve(prob::OptimizationProblem, opt::Union{Optim.Fminbox,Optim.SAMIN}, data = DEFAULT_DATA;cb = (args...) -> (false), maxiters::Number = 1000, kwargs...)
 	local x, cur, state
+
+	if data != DEFAULT_DATA
+		maxiters = length(data)
+	end
+
 	cur, state = iterate(data)
 
   	function _cb(trace)
@@ -242,6 +248,11 @@ end
 
 function __solve(prob::OptimizationProblem, opt::Optim.ConstrainedOptimizer, data = DEFAULT_DATA;cb = (args...) -> (false), maxiters::Number = 1000, kwargs...)
 	local x, cur, state
+
+	if data != DEFAULT_DATA
+		maxiters = length(data)
+	end
+
 	cur, state = iterate(data)
 
   	function _cb(trace)
@@ -309,6 +320,11 @@ function __init__()
 
 		function __solve(prob::OptimizationProblem, opt::BBO, data = DEFAULT_DATA; cb = (args...) -> (false), maxiters::Number = 1000, kwargs...)
 			local x, cur, state
+
+			if data != DEFAULT_DATA
+				maxiters = length(data)
+			end
+
 			cur, state = iterate(data)
 
 			function _cb(trace)
@@ -563,6 +579,11 @@ function __init__()
 
 		function __solve(prob::OptimizationProblem, opt::Evolutionary.AbstractOptimizer, data = DEFAULT_DATA; cb = (args...) -> (false), maxiters::Number = 1000, kwargs...)
 			local x, cur, state
+
+			if data != DEFAULT_DATA
+				maxiters = length(data)
+			end
+
 			cur, state = iterate(data)
 
 			function _cb(trace)
@@ -594,6 +615,11 @@ function __init__()
 
 		function __solve(prob::OptimizationProblem, opt::CMAEvolutionStrategyOpt, data = DEFAULT_DATA; cb = (args...) -> (false), maxiters::Number = 1000, kwargs...)
 			local x, cur, state
+
+			if data != DEFAULT_DATA
+				maxiters = length(data)
+			end
+
 			cur, state = iterate(data)
 
 			function _cb(trace)

test/minibatch.jl

@@ -1,4 +1,4 @@
-using DifferentialEquations, Flux, Optim, DiffEqFlux, Plots, GalacticOptim
+using DiffEqFlux, Plots, GalacticOptim, OrdinaryDiffEq
 
 function newtons_cooling(du, u, p, t)
     temp = u[1]
@@ -52,9 +52,10 @@ k = 10
 train_loader = Flux.Data.DataLoader(ode_data, t, batchsize = k)
 
 numEpochs = 300
+l1 = loss_adjoint(pp, train_loader.data[1], train_loader.data[2])[1]
 
 optfun = OptimizationFunction((θ, p, batch, time_batch) -> loss_adjoint(θ, batch, time_batch), GalacticOptim.AutoZygote())
 optprob = OptimizationProblem(optfun, pp)
 using IterTools: ncycle
 res1 = GalacticOptim.solve(optprob, ADAM(0.05), ncycle(train_loader, numEpochs), cb = cb, maxiters = numEpochs)
-cb(res1.minimizer, loss_adjoint(res1.minimizer, ode_data, t)...; doplot=true)
+@test 10res1.minimum < l1

test/rosenbrock.jl

@@ -1,4 +1,4 @@
-using GalacticOptim, Optim, Flux, Test
+using GalacticOptim, Optim, Test
 
 rosenbrock(x, p) =  (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
 x0 = zeros(2)

test/runtests.jl

@@ -1,4 +1,10 @@
 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