Extend train epoch schedule by warmup_epochs if warmup_prefix enabled

timm/scheduler/cosine_lr.py

@@ -111,6 +111,7 @@ def _get_lr(self, t: int) -> List[float]:
     def get_cycle_length(self, cycles=0):
         cycles = max(1, cycles or self.cycle_limit)
         if self.cycle_mul == 1.0:
-            return self.t_initial * cycles
+            t = self.t_initial * cycles
         else:
-            return int(math.floor(-self.t_initial * (self.cycle_mul ** cycles - 1) / (1 - self.cycle_mul)))
+            t = int(math.floor(-self.t_initial * (self.cycle_mul ** cycles - 1) / (1 - self.cycle_mul)))
+        return t + self.warmup_t if self.warmup_prefix else t

timm/scheduler/poly_lr.py

@@ -107,6 +107,7 @@ def _get_lr(self, t: int) -> List[float]:
     def get_cycle_length(self, cycles=0):
         cycles = max(1, cycles or self.cycle_limit)
         if self.cycle_mul == 1.0:
-            return self.t_initial * cycles
+            t = self.t_initial * cycles
         else:
-            return int(math.floor(-self.t_initial * (self.cycle_mul ** cycles - 1) / (1 - self.cycle_mul)))
+            t = int(math.floor(-self.t_initial * (self.cycle_mul ** cycles - 1) / (1 - self.cycle_mul)))
+        return t + self.warmup_t if self.warmup_prefix else t

timm/scheduler/scheduler_factory.py

@@ -196,11 +196,15 @@ def create_scheduler_v2(
         )
 
     if hasattr(lr_scheduler, 'get_cycle_length'):
-        # for cycle based schedulers (cosine, tanh, poly) recalculate total epochs w/ cycles & cooldown
+        # For cycle based schedulers (cosine, tanh, poly) recalculate total epochs w/ cycles & cooldown
+        # NOTE: Warmup prefix added in get_cycle_lengths() if enabled
         t_with_cycles_and_cooldown = lr_scheduler.get_cycle_length() + cooldown_t
         if step_on_epochs:
             num_epochs = t_with_cycles_and_cooldown
         else:
             num_epochs = t_with_cycles_and_cooldown // updates_per_epoch
+    else:
+        if warmup_prefix:
+            num_epochs += warmup_epochs
 
     return lr_scheduler, num_epochs

timm/scheduler/tanh_lr.py

@@ -108,6 +108,7 @@ def _get_lr(self, t: int) -> List[float]:
     def get_cycle_length(self, cycles=0):
         cycles = max(1, cycles or self.cycle_limit)
         if self.cycle_mul == 1.0:
-            return self.t_initial * cycles
+            t = self.t_initial * cycles
         else:
-            return int(math.floor(-self.t_initial * (self.cycle_mul ** cycles - 1) / (1 - self.cycle_mul)))
+            t = int(math.floor(-self.t_initial * (self.cycle_mul ** cycles - 1) / (1 - self.cycle_mul)))
+        return t + self.warmup_t if self.warmup_prefix else t

train.py

@@ -832,8 +832,13 @@ def main():
             lr_scheduler.step(start_epoch)
 
     if utils.is_primary(args):
+        if args.warmup_prefix:
+            sched_explain = '(warmup_epochs + epochs + cooldown_epochs). Warmup added to total when warmup_prefix=True'
+        else:
+            sched_explain = '(epochs + cooldown_epochs). Warmup within epochs when warmup_prefix=False'
         _logger.info(
-            f'Scheduled epochs: {num_epochs}. LR stepped per {"epoch" if lr_scheduler.t_in_epochs else "update"}.')
+            f'Scheduled epochs: {num_epochs} {sched_explain}. '
+            f'LR stepped per {"epoch" if lr_scheduler.t_in_epochs else "update"}.')
 
     results = []
     try: