| import math | |
| from torch.optim.lr_scheduler import LambdaLR | |
| def cosine_lr(optimizer, base_lr, warmup_length, total_steps): | |
| import math | |
| # ✅ 如果 warmup_length 是 list，则取第一个值 | |
| if isinstance(warmup_length, (list, tuple)): | |
| warmup_length = warmup_length[0] | |
| if isinstance(total_steps, (list, tuple)): | |
| total_steps = total_steps[0] | |
| def lr_lambda(current_step): | |
| if current_step < warmup_length: | |
| return float(current_step) / float(max(1, warmup_length)) | |
| progress = float(current_step - warmup_length) / float(max(1, total_steps - warmup_length)) | |
| return 0.5 * (1.0 + math.cos(math.pi * progress)) | |
| return LambdaLR(optimizer, lr_lambda) | |