ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-01-26 16:39:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Adjustments to defaults for automagic

Browse Source
This commit is contained in:
Jaret Burkett
2024-11-29 10:28:06 -07:00
parent f213996aa5
commit 92ce93140e
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
toolkit/optimizers/automagic.py
View File

@@ -13,8 +13,9 @@ class Automagic(torch.optim.Optimizer):
params,
lr=None,
min_lr=1e-7,
max_lr=1e-2,
lr_momentum=0.9,
max_lr=1e-3,
lr_pump_scale=1.1,
lr_dump_scale=0.85,
eps=(1e-30, 1e-3),
clip_threshold=1.0,
decay_rate=-0.8,
@@ -25,7 +26,8 @@ class Automagic(torch.optim.Optimizer):
self.lr = lr
self.min_lr = min_lr
self.max_lr = max_lr
self.lr_momentum = lr_momentum
self.lr_pump_scale = lr_pump_scale
self.lr_dump_scale = lr_dump_scale
defaults = {
"lr": lr,
@@ -249,8 +251,8 @@ class Automagic(torch.optim.Optimizer):
# Update learning rate mask based on sign agreement
new_lr = torch.where(
sign_agreement > 0,
lr_mask / self.lr_momentum, # Increase lr
lr_mask * self.lr_momentum # Decrease lr
lr_mask * self.lr_pump_scale, # Increase lr
lr_mask * self.lr_dump_scale # Decrease lr
)
# Clip learning rates to bounds