(current_step)
| 427 | |
| 428 | # Create a scheduler that first warms up linearly, then decays linearly |
| 429 | def lr_lambda(current_step): |
| 430 | if current_step < warmup_steps: |
| 431 | # Linear warmup from 0 to learning_rate |
| 432 | return float(current_step) / float(max(1, warmup_steps)) |
| 433 | else: |
| 434 | # Linear decay from learning_rate to 0 |
| 435 | progress = float(current_step - warmup_steps) / float( |
| 436 | max(1, max_steps - warmup_steps) |
| 437 | ) |
| 438 | return max(0.0, 1.0 - progress) |
| 439 | |
| 440 | lr_scheduler = torch.optim.lr_scheduler.LambdaLR( |
| 441 | optimizer, lr_lambda, last_epoch=-1 |
nothing calls this directly
no outgoing calls
no test coverage detected