declare-lab/tango2

1# coding=utf-82# Copyright 2023 The HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""PyTorch optimization for diffusion models."""16 17import math18from enum import Enum19from typing import Optional, Union20 21from torch.optim import Optimizer22from torch.optim.lr_scheduler import LambdaLR23 24from .utils import logging25 26 27logger = logging.get_logger(__name__)28 29 30class SchedulerType(Enum):31    LINEAR = "linear"32    COSINE = "cosine"33    COSINE_WITH_RESTARTS = "cosine_with_restarts"34    POLYNOMIAL = "polynomial"35    CONSTANT = "constant"36    CONSTANT_WITH_WARMUP = "constant_with_warmup"37 38 39def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1):40    """41    Create a schedule with a constant learning rate, using the learning rate set in optimizer.42 43    Args:44        optimizer ([`~torch.optim.Optimizer`]):45            The optimizer for which to schedule the learning rate.46        last_epoch (`int`, *optional*, defaults to -1):47            The index of the last epoch when resuming training.48 49    Return:50        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.51    """52    return LambdaLR(optimizer, lambda _: 1, last_epoch=last_epoch)53 54 55def get_constant_schedule_with_warmup(optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1):56    """57    Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate58    increases linearly between 0 and the initial lr set in the optimizer.59 60    Args:61        optimizer ([`~torch.optim.Optimizer`]):62            The optimizer for which to schedule the learning rate.63        num_warmup_steps (`int`):64            The number of steps for the warmup phase.65        last_epoch (`int`, *optional*, defaults to -1):66            The index of the last epoch when resuming training.67 68    Return:69        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.70    """71 72    def lr_lambda(current_step: int):73        if current_step < num_warmup_steps:74            return float(current_step) / float(max(1.0, num_warmup_steps))75        return 1.076 77    return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)78 79 80def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):81    """82    Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after83    a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.84 85    Args:86        optimizer ([`~torch.optim.Optimizer`]):87            The optimizer for which to schedule the learning rate.88        num_warmup_steps (`int`):89            The number of steps for the warmup phase.90        num_training_steps (`int`):91            The total number of training steps.92        last_epoch (`int`, *optional*, defaults to -1):93            The index of the last epoch when resuming training.94 95    Return:96        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.97    """98 99    def lr_lambda(current_step: int):100        if current_step < num_warmup_steps:101            return float(current_step) / float(max(1, num_warmup_steps))102        return max(103            0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))104        )105 106    return LambdaLR(optimizer, lr_lambda, last_epoch)107 108 109def get_cosine_schedule_with_warmup(110    optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1111):112    """113    Create a schedule with a learning rate that decreases following the values of the cosine function between the114    initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the115    initial lr set in the optimizer.116 117    Args:118        optimizer ([`~torch.optim.Optimizer`]):119            The optimizer for which to schedule the learning rate.120        num_warmup_steps (`int`):121            The number of steps for the warmup phase.122        num_training_steps (`int`):123            The total number of training steps.124        num_periods (`float`, *optional*, defaults to 0.5):125            The number of periods of the cosine function in a schedule (the default is to just decrease from the max126            value to 0 following a half-cosine).127        last_epoch (`int`, *optional*, defaults to -1):128            The index of the last epoch when resuming training.129 130    Return:131        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.132    """133 134    def lr_lambda(current_step):135        if current_step < num_warmup_steps:136            return float(current_step) / float(max(1, num_warmup_steps))137        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))138        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))139 140    return LambdaLR(optimizer, lr_lambda, last_epoch)141 142 143def get_cosine_with_hard_restarts_schedule_with_warmup(144    optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1145):146    """147    Create a schedule with a learning rate that decreases following the values of the cosine function between the148    initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases149    linearly between 0 and the initial lr set in the optimizer.150 151    Args:152        optimizer ([`~torch.optim.Optimizer`]):153            The optimizer for which to schedule the learning rate.154        num_warmup_steps (`int`):155            The number of steps for the warmup phase.156        num_training_steps (`int`):157            The total number of training steps.158        num_cycles (`int`, *optional*, defaults to 1):159            The number of hard restarts to use.160        last_epoch (`int`, *optional*, defaults to -1):161            The index of the last epoch when resuming training.162 163    Return:164        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.165    """166 167    def lr_lambda(current_step):168        if current_step < num_warmup_steps:169            return float(current_step) / float(max(1, num_warmup_steps))170        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))171        if progress >= 1.0:172            return 0.0173        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0))))174 175    return LambdaLR(optimizer, lr_lambda, last_epoch)176 177 178def get_polynomial_decay_schedule_with_warmup(179    optimizer, num_warmup_steps, num_training_steps, lr_end=1e-7, power=1.0, last_epoch=-1180):181    """182    Create a schedule with a learning rate that decreases as a polynomial decay from the initial lr set in the183    optimizer to end lr defined by *lr_end*, after a warmup period during which it increases linearly from 0 to the184    initial lr set in the optimizer.185 186    Args:187        optimizer ([`~torch.optim.Optimizer`]):188            The optimizer for which to schedule the learning rate.189        num_warmup_steps (`int`):190            The number of steps for the warmup phase.191        num_training_steps (`int`):192            The total number of training steps.193        lr_end (`float`, *optional*, defaults to 1e-7):194            The end LR.195        power (`float`, *optional*, defaults to 1.0):196            Power factor.197        last_epoch (`int`, *optional*, defaults to -1):198            The index of the last epoch when resuming training.199 200    Note: *power* defaults to 1.0 as in the fairseq implementation, which in turn is based on the original BERT201    implementation at202    https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37203 204    Return:205        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.206 207    """208 209    lr_init = optimizer.defaults["lr"]210    if not (lr_init > lr_end):211        raise ValueError(f"lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})")212 213    def lr_lambda(current_step: int):214        if current_step < num_warmup_steps:215            return float(current_step) / float(max(1, num_warmup_steps))216        elif current_step > num_training_steps:217            return lr_end / lr_init  # as LambdaLR multiplies by lr_init218        else:219            lr_range = lr_init - lr_end220            decay_steps = num_training_steps - num_warmup_steps221            pct_remaining = 1 - (current_step - num_warmup_steps) / decay_steps222            decay = lr_range * pct_remaining**power + lr_end223            return decay / lr_init  # as LambdaLR multiplies by lr_init224 225    return LambdaLR(optimizer, lr_lambda, last_epoch)226 227 228TYPE_TO_SCHEDULER_FUNCTION = {229    SchedulerType.LINEAR: get_linear_schedule_with_warmup,230    SchedulerType.COSINE: get_cosine_schedule_with_warmup,231    SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup,232    SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup,233    SchedulerType.CONSTANT: get_constant_schedule,234    SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup,235}236 237 238def get_scheduler(239    name: Union[str, SchedulerType],240    optimizer: Optimizer,241    num_warmup_steps: Optional[int] = None,242    num_training_steps: Optional[int] = None,243    num_cycles: int = 1,244    power: float = 1.0,245    last_epoch: int = -1,246):247    """248    Unified API to get any scheduler from its name.249 250    Args:251        name (`str` or `SchedulerType`):252            The name of the scheduler to use.253        optimizer (`torch.optim.Optimizer`):254            The optimizer that will be used during training.255        num_warmup_steps (`int`, *optional*):256            The number of warmup steps to do. This is not required by all schedulers (hence the argument being257            optional), the function will raise an error if it's unset and the scheduler type requires it.258        num_training_steps (`int``, *optional*):259            The number of training steps to do. This is not required by all schedulers (hence the argument being260            optional), the function will raise an error if it's unset and the scheduler type requires it.261        num_cycles (`int`, *optional*):262            The number of hard restarts used in `COSINE_WITH_RESTARTS` scheduler.263        power (`float`, *optional*, defaults to 1.0):264            Power factor. See `POLYNOMIAL` scheduler265        last_epoch (`int`, *optional*, defaults to -1):266            The index of the last epoch when resuming training.267    """268    name = SchedulerType(name)269    schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name]270    if name == SchedulerType.CONSTANT:271        return schedule_func(optimizer, last_epoch=last_epoch)272 273    # All other schedulers require `num_warmup_steps`274    if num_warmup_steps is None:275        raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.")276 277    if name == SchedulerType.CONSTANT_WITH_WARMUP:278        return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, last_epoch=last_epoch)279 280    # All other schedulers require `num_training_steps`281    if num_training_steps is None:282        raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.")283 284    if name == SchedulerType.COSINE_WITH_RESTARTS:285        return schedule_func(286            optimizer,287            num_warmup_steps=num_warmup_steps,288            num_training_steps=num_training_steps,289            num_cycles=num_cycles,290            last_epoch=last_epoch,291        )292 293    if name == SchedulerType.POLYNOMIAL:294        return schedule_func(295            optimizer,296            num_warmup_steps=num_warmup_steps,297            num_training_steps=num_training_steps,298            power=power,299            last_epoch=last_epoch,300        )301 302    return schedule_func(303        optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, last_epoch=last_epoch304    )305