wehub-resource-sync
106 lines
3.8 KiB
Python
Executable File
106 lines
3.8 KiB
Python
Executable File
# 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)
|