Provide default progress loggers and use new ProgressLogging API

Project.toml

@@ -10,10 +10,10 @@ DiffEqSensitivity = "41bf760c-e81c-5289-8e54-58b1f1f8abe2"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-Juno = "e5e0dc1b-0480-54bc-9374-aad01c23163d"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
@@ -28,21 +28,20 @@ DiffEqSensitivity = "6.7"
 DiffResults = "0.0.4, 1.0"
 Flux = "0.10.1"
 ForwardDiff = "0.10"
-Juno = "0.7, 0.8"
 Optim = "0.20"
+ProgressLogging = "0.1"
 RecursiveArrayTools = "2"
 Requires = "0.5, 1.0"
 ReverseDiff = "1.1"
 StaticArrays = "0.11, 0.12"
 Tracker = "0.2"
-Zygote = "0.4"
 ZygoteRules = "0.2"
 julia = "1"
 
 [extras]
 DelayDiffEq = "bcd4f6db-9728-5f36-b5f7-82caef46ccdb"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 NLopt = "76087f3c-5699-56af-9a33-bf431cd00edd"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"

src/DiffEqFlux.jl

@@ -1,10 +1,12 @@
 module DiffEqFlux
 
 using DiffEqBase, Tracker, DiffResults, DiffEqSensitivity, ForwardDiff,
-      Flux, Requires, Adapt, LinearAlgebra, RecursiveArrayTools, Juno, Optim,
-      StaticArrays, Base.Iterators, Logging
+      Flux, Requires, Adapt, LinearAlgebra, RecursiveArrayTools, Optim,
+      StaticArrays, Base.Iterators
 
-import ZygoteRules, ReverseDiff
+import ProgressLogging, ZygoteRules, ReverseDiff
+
+import Logging
 
 gpu_or_cpu(x) = Array
 

src/train.jl

@@ -36,6 +36,17 @@ struct TrackerDiffMode <: DiffMode end
 struct ReverseDiffMode <: DiffMode end
 struct ZygoteDiffMode <: DiffMode end
 
+macro withprogress(progress, exprs...)
+  quote
+    if $progress
+      $DiffEqBase.maybe_with_logger($DiffEqBase.default_logger($Logging.current_logger())) do
+        $ProgressLogging.@withprogress $(exprs...)
+      end
+    else
+      $(exprs[end])
+    end
+  end |> esc
+end
 
 """
     sciml_train(loss, θ, opt, data = DEFAULT_DATA; cb, maxiters)
@@ -55,8 +66,8 @@ The keyword arguments are as follows:
   defaults to infinite.
 """
 function sciml_train(loss, _θ, opt, _data = DEFAULT_DATA;
-                     cb = (args...) -> (false),
-                     maxiters = get_maxiters(data),progress=true)
+                     cb = (args...) -> false,
+                     maxiters = get_maxiters(data), progress=true)
 
   # Flux is silly and doesn't have an abstract type on its optimizers, so assume
   # this is a Flux optimizer
@@ -73,29 +84,24 @@ function sciml_train(loss, _θ, opt, _data = DEFAULT_DATA;
     data = _data
   end
 
-  progress && @logmsg(LogLevel(-1),
-  "Training",
-  _id = :DiffEqFlux,
-  progress=0)
-
   t0 = time()
+
   local x
-  @progress for (i,d) in enumerate(data)
-    gs = Flux.Zygote.gradient(ps) do
-      x = loss(θ,d...)
-      first(x)
-    end
-    cb_call = cb(θ,x...)
-    if !(typeof(cb_call) <: Bool)
-      error("The callback should return a boolean `halt` for whether to stop the optimization process. Please see the sciml_train documentation for information.")
-    elseif cb_call
-      break
+  @withprogress progress name="Training" begin
+    for (i,d) in enumerate(data)
+      gs = Flux.Zygote.gradient(ps) do
+        x = loss(θ,d...)
+        first(x)
+      end
+      cb_call = cb(θ,x...)
+      if !(typeof(cb_call) <: Bool)
+        error("The callback should return a boolean `halt` for whether to stop the optimization process. Please see the sciml_train documentation for information.")
+      elseif cb_call
+        break
+      end
+      progress && ProgressLogging.@logprogress i/maxiters
+      update!(opt, ps, gs)
     end
-    progress && @logmsg(LogLevel(-1),
-    "Training",
-    _id = :DiffEqFlux,
-    progress=i/maxiters)
-    update!(opt, ps, gs)
   end
 
   _time = time()
@@ -126,7 +132,7 @@ function sciml_train(loss, _θ, opt, _data = DEFAULT_DATA;
                                         0,
                                         true,
                                         NaN,
-                                        _time-t0,)
+                                        _time-t0)
 end
 
 decompose_trace(trace::Optim.OptimizationTrace) = last(trace)