VT for manoptuna

JuliaManifolds · Jan 3, 2024 · 6e334bc · 6e334bc
1 parent f7e1435
commit 6e334bc
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 5 deletions.
diff --git a/benchmarks/benchmark_comparison.jl b/benchmarks/benchmark_comparison.jl
@@ -149,10 +149,12 @@ function benchmark_time_state(
     x0,
     stepsize,
     mem_len::Int,
-    gtol::Real,
+    gtol::Real;
+    kwargs...,
 )
     manopt_sc = StopWhenGradientInfNormLess(gtol) | StopAfterIteration(1000)
     M = manifold_maker(manifold_name, N, :Manopt)
+    mem_len = min(mem_len, manifold_dimension(M))
     bench_manopt = @benchmark quasi_Newton(
         $M,
         $f_manopt,
@@ -162,6 +164,7 @@ function benchmark_time_state(
         evaluation=$(InplaceEvaluation()),
         memory_size=$mem_len,
         stopping_criterion=$(manopt_sc),
+        $kwargs...,
     )
     manopt_state = quasi_Newton(
         M,
@@ -173,6 +176,7 @@ function benchmark_time_state(
         return_state=true,
         memory_size=mem_len,
         stopping_criterion=manopt_sc,
+        kwargs...,
     )
     iters = get_count(manopt_state, :Iterations)
     final_val = f_manopt(M, manopt_state.p)
@@ -184,6 +188,7 @@ struct OptimQN <: AbstractOptimConfig end
 function benchmark_time_state(
     ::OptimQN, manifold_name, N, f, g!, x0, stepsize, mem_len::Int, gtol::Real
 )
+    mem_len = min(mem_len, manifold_dimension(manifold_maker(manifold_name, N, :Manopt)))
     options_optim = Optim.Options(; g_tol=gtol)
     method_optim = LBFGS(;
         m=mem_len, linesearch=stepsize, manifold=manifold_maker(manifold_name, N, :Optim)

diff --git a/benchmarks/manoptuna.jl b/benchmarks/manoptuna.jl
@@ -1,4 +1,4 @@
-
+using Manifolds
 using PythonCall
 include("benchmark_comparison.jl")
 
@@ -58,17 +58,25 @@ function lbfgs_objective(trial)
     ls_hz = LineSearches.HagerZhang()
 
     N_range = [2^n for n in 1:3:16]
+    vts = [ParallelTransport(), ProjectionTransport()]
+    vt = vts[pyconvert(Int, trial.suggest_categorical("vector_transport_method", (1, 2)))]
 
     # TODO: ensure this actually somewhat realistic,
     # otherwise there is too little pruning (if values here are too low)
     # or too much pruning (if values here are too high)
     # regenerate using
     # prunining_losses = lbfgs_compute_pruning_losses()
     # *but* with zeroed-out prunning_losses
-    prunining_losses = [56.403, 69.438, 96.36449999999999, 409.749, 2542.482, 6.366860307e6]
+    # padded with zeros for convenience
+    prunining_losses = vcat(
+        [56.403, 69.438, 96.36449999999999, 409.749, 2542.482, 6.366860307e6], zeros(100)
+    )
 
     loss = sum(prunining_losses)
 
+    # here iterate over problems we want to optimize for
+    # from smallest to largest; pruning should stop the iteration early
+    # if the hyperparameter set is not promising
     cur_i = 0
     for N in N_range
         x0 = zeros(N)
@@ -82,8 +90,10 @@ function lbfgs_objective(trial)
             x0,
             Manopt.LineSearchesStepsize(ls_hz),
             pyconvert(Int, mem_len),
-            gtol,
+            gtol;
+            vector_transport_method=vt,
         )
+        # TODO: take objective_value into account for loss?
         loss -= prunining_losses[cur_i + 1]
         loss += manopt_time
         trial.report(loss, cur_i)
@@ -97,6 +107,6 @@ end
 
 function lbfgs_study()
     study = optuna.create_study(; study_name="L-BFGS")
-    study.optimize(lbfgs_objective; n_trials=1000, timeout=200)
+    study.optimize(lbfgs_objective; n_trials=1000, timeout=500)
     return println("Best params is $(study.best_params) with value $(study.best_value)")
 end