File size: 721 Bytes
b91f38a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | 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)
|