42 lines
812 B
Julia
42 lines
812 B
Julia
using Juno: info
|
|
|
|
"""
|
|
@cb for ... end t expr
|
|
|
|
Run the for loop, executing `expr` every `t` seconds.
|
|
"""
|
|
macro cb(ex, t, f)
|
|
@assert isexpr(ex, :for)
|
|
cond, body = ex.args
|
|
@esc t f cond body
|
|
:(let
|
|
t0 = time_ns()
|
|
dt = $t*1e9
|
|
@progress $(Expr(:for, cond, quote
|
|
t = time_ns()
|
|
if t - t0 > dt
|
|
t0 = t
|
|
f = () -> $f
|
|
f()
|
|
end
|
|
$body
|
|
end))
|
|
end)
|
|
end
|
|
|
|
function train!(m, train; cb = [],
|
|
epoch = 1, η = 0.1, loss = mse)
|
|
@progress for e in 1:epoch
|
|
info("Epoch $e")
|
|
@cb for (x, y) in train
|
|
x, y = tobatch.((x, y))
|
|
ŷ = m(x)
|
|
any(isnan, ŷ) && error("NaN")
|
|
Δ = back!(loss, 1, ŷ, y)
|
|
back!(m, Δ, x)
|
|
update!(m, η)
|
|
end 5 foreach(f -> f(), cb)
|
|
end
|
|
return m
|
|
end
|