Spaces:
Runtime error
Runtime error
File size: 3,862 Bytes
f1f9265 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0
import math
from diffusers import get_constant_schedule_with_warmup, get_cosine_schedule_with_warmup
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR
from diffusion.utils.logger import get_root_logger
def build_lr_scheduler(config, optimizer, train_dataloader, lr_scale_ratio):
if not config.get("lr_schedule_args", None):
config.lr_schedule_args = dict()
if config.get("lr_warmup_steps", None):
config["num_warmup_steps"] = config.get("lr_warmup_steps") # for compatibility with old version
logger = get_root_logger()
logger.info(
f"Lr schedule: {config.lr_schedule}, "
+ ",".join([f"{key}:{value}" for key, value in config.lr_schedule_args.items()])
+ "."
)
if config.lr_schedule == "cosine":
lr_scheduler = get_cosine_schedule_with_warmup(
optimizer=optimizer,
**config.lr_schedule_args,
num_training_steps=(len(train_dataloader) * config.num_epochs),
)
elif config.lr_schedule == "constant":
lr_scheduler = get_constant_schedule_with_warmup(
optimizer=optimizer,
**config.lr_schedule_args,
)
elif config.lr_schedule == "cosine_decay_to_constant":
assert lr_scale_ratio >= 1
lr_scheduler = get_cosine_decay_to_constant_with_warmup(
optimizer=optimizer,
**config.lr_schedule_args,
final_lr=1 / lr_scale_ratio,
num_training_steps=(len(train_dataloader) * config.num_epochs),
)
else:
raise RuntimeError(f"Unrecognized lr schedule {config.lr_schedule}.")
return lr_scheduler
def get_cosine_decay_to_constant_with_warmup(
optimizer: Optimizer,
num_warmup_steps: int,
num_training_steps: int,
final_lr: float = 0.0,
num_decay: float = 0.667,
num_cycles: float = 0.5,
last_epoch: int = -1,
):
"""
Create a schedule with a cosine annealing lr followed by a constant lr.
Args:
optimizer ([`~torch.optim.Optimizer`]):
The optimizer for which to schedule the learning rate.
num_warmup_steps (`int`):
The number of steps for the warmup phase.
num_training_steps (`int`):
The number of total training steps.
final_lr (`int`):
The final constant lr after cosine decay.
num_decay (`int`):
The
last_epoch (`int`, *optional*, defaults to -1):
The index of the last epoch when resuming training.
Return:
`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
"""
def lr_lambda(current_step):
if current_step < num_warmup_steps:
return float(current_step) / float(max(1, num_warmup_steps))
num_decay_steps = int(num_training_steps * num_decay)
if current_step > num_decay_steps:
return final_lr
progress = float(current_step - num_warmup_steps) / float(max(1, num_decay_steps - num_warmup_steps))
return (
max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))) * (1 - final_lr) + final_lr
)
return LambdaLR(optimizer, lr_lambda, last_epoch)
|