Get learning rate
(it, training_params: TrainingParams)
| 777 | |
| 778 | # learning rate decay scheduler (cosine with warmup) |
| 779 | def get_lr(it, training_params: TrainingParams): |
| 780 | """Get learning rate""" |
| 781 | |
| 782 | # 1) linear warmup for warmup_iters steps |
| 783 | if it < training_params['warmup_iters']: |
| 784 | return training_params['learning_rate'] * it / training_params['warmup_iters'] |
| 785 | |
| 786 | # 2) if it > lr_decay_iters, return min learning rate |
| 787 | if it > training_params['lr_decay_iters']: |
| 788 | return training_params['min_lr'] |
| 789 | |
| 790 | # 3) in between, use cosine decay down to min learning rate |
| 791 | decay_ratio = (it - training_params['warmup_iters']) / ( |
| 792 | training_params['lr_decay_iters'] - training_params['warmup_iters'] |
| 793 | ) |
| 794 | assert 0 <= decay_ratio <= 1 |
| 795 | coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # coeff ranges 0..1 |
| 796 | return training_params['min_lr'] + coeff * ( |
| 797 | training_params['learning_rate'] - training_params['min_lr'] |
| 798 | ) |
| 799 | |
| 800 | |
| 801 | def save_checkpoint( |