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

fix: scheduler division-by-zero, loss_mask bool

Browse Source 
- schedule.py: guard warmup_steps/lr_decay_steps against zero
- strategy.py: use ~loss_mask instead of loss_mask==0 on bool tensor
This commit is contained in:
ViperEkura
2026-07-03 22:04:55 +08:00
parent 70c0e5de90
commit 27d1921d9c
2 changed files with 10 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/trainer/schedule.py
+9 -3
View File
@@ -65,11 +65,15 @@ class CosineScheduler(BaseScheduler):
def get_lr(self) -> List[float]:
# warmup
if self.last_epoch < self.warmup_steps:
warmup_factor = max(self.min_rate, self.last_epoch / self.warmup_steps)
warmup_factor = max(
self.min_rate, self.last_epoch / max(self.warmup_steps, 1)
)
return [base_lr * warmup_factor for base_lr in self.base_lrs]
# cosine decay
decay_progress = (self.last_epoch - self.warmup_steps) / self.lr_decay_steps
decay_progress = (self.last_epoch - self.warmup_steps) / max(
self.lr_decay_steps, 1
)
decay_progress = min(decay_progress, 1.0)
cosine_decay = 0.5 * (1.0 + math.cos(math.pi * decay_progress))
decay_factor = max(self.min_rate, cosine_decay)
@@ -118,7 +122,9 @@ class SGDRScheduler(BaseScheduler):
def get_lr(self):
# warmup
if self.last_epoch < self.warmup_steps:
warmup_factor = max(self.min_rate, self.last_epoch / self.warmup_steps)
warmup_factor = max(
self.min_rate, self.last_epoch / max(self.warmup_steps, 1)
)
return [base_lr * warmup_factor for base_lr in self.base_lrs]
# SGDR