ViperEkura
/
AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: unify scheduler min_rate default to 0.01, clamp WSD warmup

This commit is contained in:
ViperEkura 2026-07-03 17:52:23 +08:00
parent dc9faca3b1
commit 500c605fad
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
astrai/trainer/schedule.py
View File

@ -53,7 +53,7 @@ class CosineScheduler(BaseScheduler):
optimizer,
warmup_steps: int,
lr_decay_steps: int,
min_rate: float = 0.05,
min_rate: float = 0.01,
last_epoch: int = -1,
):
self.warmup_steps = warmup_steps
@ -104,7 +104,7 @@ class SGDRScheduler(BaseScheduler):
optimizer,
warmup_steps: int,
cycle_length: int,
min_rate: float = 0.05,
min_rate: float = 0.01,
t_mult: int = 2,
last_epoch: int = -1,
):
@ -182,7 +182,7 @@ class WSDScheduler(BaseScheduler):
warmup_steps: int,
stable_steps: int,
decay_steps: int,
min_rate: float = 0.0,
min_rate: float = 0.01,
last_epoch: int = -1,
):
self.warmup_steps = warmup_steps
@ -194,7 +194,7 @@ class WSDScheduler(BaseScheduler):
def get_lr(self) -> List[float]:
if self.last_epoch < self.warmup_steps:
factor = self.last_epoch / max(self.warmup_steps, 1)
factor = max(self.min_rate, self.last_epoch / max(self.warmup_steps, 1))
return [base_lr * factor for base_lr in self.base_lrs]
offset = self.last_epoch - self.warmup_steps