fix nested breaking and iteration counts

Author: Vaibhavdixit02

lib/OptimizationOptimisers/src/OptimizationOptimisers.jl

@@ -79,8 +79,12 @@ function SciMLBase.__solve(cache::OptimizationCache{
     fevals = 0
     gevals = 0
     t0 = time()
+    breakall = false
     Optimization.@withprogress cache.progress name="Training" begin
         for epoch in 1:epochs
+            if breakall
+                break
+            end
             for (i, d) in enumerate(data)
                 if cache.f.fg !== nothing && dataiterate
                     x = cache.f.fg(G, θ, d)
@@ -111,10 +115,10 @@ function SciMLBase.__solve(cache::OptimizationCache{
                     objective = x[1],
                     grad = G,
                     original = state)
-                cb_call = cache.callback(opt_state, x...)
-                if !(cb_call isa Bool)
+                breakall = cache.callback(opt_state, x...)
+                if !(breakall isa Bool)
                     error("The callback should return a boolean `halt` for whether to stop the optimization process. Please see the `solve` documentation for information.")
-                elseif cb_call
+                elseif breakall
                     break
                 end
                 msg = @sprintf("loss: %.3g", first(x)[1])
@@ -126,7 +130,7 @@ function SciMLBase.__solve(cache::OptimizationCache{
                         min_err = x
                         min_θ = copy(θ)
                     end
-                    if i == length(data)  #Last iter, revert to best.
+                    if i == length(data)*epochs  #Last iter, revert to best.
                         opt = min_opt
                         x = min_err
                         θ = min_θ
@@ -136,7 +140,7 @@ function SciMLBase.__solve(cache::OptimizationCache{
                             objective = x[1],
                             grad = G,
                             original = state)
-                        cache.callback(opt_state, x...)
+                        breakall = cache.callback(opt_state, x...)
                         break
                     end
                 end

lib/OptimizationOptimisers/test/runtests.jl

@@ -76,7 +76,7 @@ end
     using Optimization, OptimizationOptimisers, Lux, Zygote, MLUtils, Random,
           ComponentArrays
 
-    x = rand(10000)
+    x = rand(Float32, 10000)
     y = sin.(x)
     data = MLUtils.DataLoader((x, y), batchsize = 100)
 
@@ -99,13 +99,16 @@ end
     optf = OptimizationFunction(loss, AutoZygote())
     prob = OptimizationProblem(optf, ps_ca, data)
 
-    res = Optimization.solve(prob, Optimisers.Adam(), callback = callback, epochs = 10000)
+    res = Optimization.solve(prob, Optimisers.Adam(), epochs = 50)
 
     @test res.objective < 1e-4
-    @test res.stats.iterations == 10000*length(data)
-    @test res.stats.fevals == 10000*length(data)
-    @test res.stats.gevals == 10000*length(data)
+    @test res.stats.iterations == 50*length(data)
+    @test res.stats.fevals == 50*length(data)
+    @test res.stats.gevals == 50*length(data)
+
+    res = Optimization.solve(prob, Optimisers.Adam(), callback = callback, epochs = 100)
 
+    @test res.objective < 1e-4
 
     using MLDataDevices
     data = CPUDevice()(data)