Fix lr scheduler when accumulating gradients

hmorimitsu · Jun 12, 2024 · 25295e8 · 25295e8
1 parent d0291f4
commit 25295e8
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/ptlflow/models/base_model/base_model.py b/ptlflow/models/base_model/base_model.py
@@ -552,10 +552,16 @@ def configure_optimizers(self) -> Dict[str, Any]:
         optimizer = optim.AdamW(
             self.parameters(), lr=self.args.lr, weight_decay=self.args.wdecay
         )
+
+        accumulate_grad_batches = (
+            1
+            if self.args.accumulate_grad_batches is None
+            else self.args.accumulate_grad_batches
+        )
         lr_scheduler = optim.lr_scheduler.OneCycleLR(
             optimizer,
             self.args.lr,
-            total_steps=self.args.max_steps,
+            total_steps=self.args.max_steps // accumulate_grad_batches,
             pct_start=0.05,
             cycle_momentum=False,
             anneal_strategy="linear",