Spaces:
Running on Zero
Running on Zero
| import math | |
| import torch.optim as optim | |
| def cosine_with_warmup(step, *, warmup_steps, total_steps, min_lr=0.0): | |
| if step < warmup_steps: | |
| return step / max(1, warmup_steps) | |
| # calculate step within warmup | |
| progress = (step - warmup_steps) / max(1, total_steps - warmup_steps) | |
| # calculate where on cosine curve we are | |
| cosine = 0.5 * (1 + math.cos(math.pi * progress)) | |
| return min_lr + (1 - min_lr) * cosine | |
| def create_cosine_scheduler(optimizer, total_steps, warmup_fraction=0.05, min_lr_fraction=0.01): | |
| warmup_steps = int(total_steps * warmup_fraction) | |
| return optim.lr_scheduler.LambdaLR( | |
| optimizer, | |
| lr_lambda=lambda s: cosine_with_warmup( | |
| s, | |
| warmup_steps=warmup_steps, | |
| total_steps=total_steps, | |
| min_lr=min_lr_fraction | |
| ) | |
| ) |