Merge pull request #1859 from darsnack/rm-juno

Use ProgressLogging instead of Juno

Author: darsnack

Project.toml

@@ -11,13 +11,13 @@ CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
-Juno = "e5e0dc1b-0480-54bc-9374-aad01c23163d"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SHA = "ea8e919c-243c-51af-8825-aaa63cd721ce"
@@ -36,10 +36,10 @@ CUDA = "3"
 CodecZlib = "0.7"
 Colors = "0.12"
 Functors = "0.2.1"
-Juno = "0.8"
 MacroTools = "0.5"
 NNlib = "0.8"
 NNlibCUDA = "0.2"
+ProgressLogging = "0.1"
 Reexport = "0.2, 1.0"
 StatsBase = "0.33"
 ZipFile = "0.9"

src/Flux.jl

@@ -4,7 +4,7 @@ module Flux
 
 using Base: tail
 using Statistics, Random, LinearAlgebra
-using Zygote, MacroTools, Juno, Reexport
+using Zygote, MacroTools, ProgressLogging, Reexport
 using MacroTools: @forward
 @reexport using NNlib
 using Zygote: Params, @adjoint, gradient, pullback, @nograd

src/data/tree.jl

@@ -21,13 +21,6 @@ function Base.show(io::IO, t::Tree)
   print_tree(io, t)
 end
 
-using Juno
-
-@render Juno.Inline t::Tree begin
-  render(t) = Juno.Tree(t.value, render.(t.children))
-  Juno.Tree(typeof(t), [render(t)])
-end
-
 Base.getindex(t::Tree, i::Integer) = t.children[i]
 Base.getindex(t::Tree, i::Integer, is::Integer...) = t[i][is...]
 

src/optimise/train.jl

@@ -1,4 +1,4 @@
-using Juno
+using ProgressLogging: @progress, @withprogress, @logprogress
 import Zygote: Params, gradient
 
 """
@@ -104,7 +104,8 @@ Multiple optimisers and callbacks can be passed to `opt` and `cb` as arrays.
 function train!(loss, ps, data, opt; cb = () -> ())
   ps = Params(ps)
   cb = runall(cb)
-  @progress for d in data
+  n = (Base.IteratorSize(typeof(data)) == Base.HasLength()) ? length(data) : 0
+  @withprogress for (i, d) in enumerate(data)
     try
       gs = gradient(ps) do
         loss(batchmemaybe(d)...)
@@ -120,6 +121,7 @@ function train!(loss, ps, data, opt; cb = () -> ())
         rethrow(ex)
       end
     end
+    @logprogress i / n
   end
 end