Optimize for N epochs. Invoke custom callback at the end of an epoch.

I want to propose some changes but I am not entirely sure where they would fit best.

Let's assume there is an iterator for a training set. Often, the suggested method to train/optimize for N epochs is to ncycle the training set. I can define a callback to log some information (e.g. the loss) which gets passed to the optimizer and invoked at the end of an iteration.
Because the training set can contain many samples I don't want to print the loss of every sample but the mean at the end of an epoch (all samples have been taken). Printing the samples' losses should still be possible though.

I can imagine a __solve() function accepting all the optimizer and some additional arguments (epochs, epoch_cb).

Here are some parts of my code. It already works but could be more customizable.

# maybe something like __solve(prob::OptimizationProblem, opt, data;
                 cb = (args...) -> (false), ecb = (args...) -> (false), epochs = nothing,
                 kwargs...)

  porgress = ProgressMeter.Progress(length(data); showspeed=true)
  losses = Float32[]
  function _cb(p,l,args...; kwargs...)
    push!(losses,l)
    x = "$(losses[end])"
    ProgressMeter.next!(porgress; showvalues = [(:loss,x)])
    cb(p,l,args...; kwargs...)
  end
  ...
  optprob = GalacticOptim.OptimizationProblem(...)
  ...
  porgress = ProgressMeter.Progress(length(data); showspeed=true)
  res = GalacticOptim.solve(optprob, opt, data, cb = _cb)
  ecb(losses,1,res)
  losses = Float32[]
  for epoch in 2:epochs
    prob = remake(optprob,u0=res.u)
    porgress = ProgressMeter.Progress(length(data); showspeed=true)
    res = GalacticOptim.solve(prob, opt, data, cb = _cb)
    ecb(losses,epoch,res)
    losses = Float32[]
  end
  return res
end

The output I get:

Epoch 019/100, mean train_loss:3.6367512
Progress: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████| Time: 0:00:44 ( 1.17  s/it)
  loss:  3.7939742
Epoch 020/100, mean train_loss:3.578114
Progress:  61%|███████████████████████████████████████████████████████████████                                         |  ETA: 0:00:18 ( 1.20  s/it)
  loss:  2.0184839

Is something like this possible already? If not and you think it's worth creating a PR I am happy to do so.

Start by examining the solve API and its existing callback handling, then compare the requested epoch and sample-level behavior described in the issue. Done would mean a clearly defined way to run for N epochs and invoke a custom callback at each epoch while preserving per-sample callbacks.