INNER CODE UNIT · Python
compute_learning_rate
OpenMeshLab/MeshXL · engine.py:10
def compute_learning_rate(args, curr_iter, max_iters):
assert curr_iter <= max_iters and curr_iter >= 0
if (curr_iter <= args.warm_lr_iters) and args.warm_lr_iters > 0:
# Linear Warmup: warm_lr -> curr_lr -> base_lr
curr_lr = args.warm_lr + curr_iter / args.warm_lr_iters * (args.base_lr - args.warm_lr)
else:
# Cosine Learning Rate Schedule
curr_lr = args.final_lr + 0.5 * (args.base_lr - args.final_lr) * (
1 + math.cos(math.pi * curr_iter / max_iters)
)
return curr_lr
def adjust_learning_rate(args, optimizer, curr_iter, max_iters):
curr_lr = compute_learning_rate(args, curr_iter, max_iters)
for param_group in optimizer.param_groups:
param_group["lr"] = curr_lr