Merge pull request #354 from kozistr/fix/schedulefreeradam-optimizer

[Fix] bias correction2 in ScheduleFreeRAdam optimizer
kozistr · Feb 23, 2025 · 3d20627 · 3d20627
2 parents c09d18b + ad60eb0
commit 3d20627
Show file tree

Hide file tree

Showing 6 changed files with 13 additions and 4 deletions.
diff --git a/docs/changelogs/v3.4.3.md b/docs/changelogs/v3.4.3.md
@@ -0,0 +1,6 @@
+### Change Log
+
+
+### Fix
+
+* bias_correction2 in ScheduleFreeRAdam optimizer. (#354)
diff --git a/docs/visualizations/rastrigin_ScheduleFreeRAdam.png b/docs/visualizations/rastrigin_ScheduleFreeRAdam.png
diff --git a/docs/visualizations/rosenbrock_ScheduleFreeRAdam.png b/docs/visualizations/rosenbrock_ScheduleFreeRAdam.png
diff --git a/examples/visualize_optimizers.py b/examples/visualize_optimizers.py
@@ -37,8 +37,8 @@
     'dadaptlion': {'lr': hp.uniform('lr', 0, 10)},
     'padam': {'lr': hp.uniform('lr', 0, 10)},
     'dadaptadam': {'lr': hp.uniform('lr', 0, 10)},
-    'adahessian': {'lr': hp.uniform('lr', 0, 800)},  # Wider range for second-order optimizers
-    'sophiah': {'lr': hp.uniform('lr', 0, 60)},  # Wider range for second-order optimizers
+    'adahessian': {'lr': hp.uniform('lr', 0, 800)},
+    'sophiah': {'lr': hp.uniform('lr', 0, 60)},
     'pid': {
         'lr': hp.uniform('lr', 0, 0.5),
         'derivative': hp.quniform('derivative', 2, 14, 0.5),
@@ -86,6 +86,9 @@
         'lr': hp.uniform('lr', 0, 3),
         'momentum': hp.quniform('momentum', 0, 0.99, 0.01),
     },
+    'schedulefreeradam': {
+        'lr': hp.uniform('lr', 1, 10),
+    },
     'kron': {
         'lr': hp.uniform('lr', 0, 0.8),
         'momentum': hp.quniform('momentum', 0, 0.99, 0.01),

diff --git a/pytorch_optimizer/optimizer/schedulefree.py b/pytorch_optimizer/optimizer/schedulefree.py
@@ -396,7 +396,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
             beta1, beta2 = group['betas']
 
-            bias_correction2: float = self.debias_beta(beta2, group['step'])
+            bias_correction2: float = self.debias(beta2, group['step'])
 
             lr, n_sma = self.get_rectify_step_size(
                 is_rectify=True,

diff --git a/tests/constants.py b/tests/constants.py
@@ -503,7 +503,7 @@
     (Adalite, {'lr': 1e0, 'weight_decay': 1e-3}, 5),
     (ScheduleFreeSGD, {'lr': 1e0, 'weight_decay': 1e-3}, 5),
     (ScheduleFreeAdamW, {'lr': 1e0, 'weight_decay': 1e-3}, 5),
-    (ScheduleFreeRAdam, {'lr': 1e0}, 20),
+    (ScheduleFreeRAdam, {'lr': 1e2, 'weight_decay': 1e-3}, 20),
     (FAdam, {'lr': 1e0, 'weight_decay': 1e-3}, 5),
     (GrokFastAdamW, {'lr': 5e0, 'weight_decay': 1e-3, 'grokfast_after_step': 1}, 5),
     (Kate, {'lr': 5e-2}, 10),