#!/usr/bin/env python3
# -*- coding: utf-8 -*-


# author: <NAME>
# file: lr_scheduler.py
# PyTorch optimization for BERT model.


import torch
import math
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR



def build_lr_scheduler(config, optimizer, total_train_step):
    lr_scheduler_type = config.lr_scheduler_type
    if lr_scheduler_type == "polynomial_warmup":
        scheduler = get_polynomial_decay_schedule_with_warmup(optimizer, config.warmup_steps, total_train_step,
                                                              lr_end=config.lr_min, power=1.0, last_epoch=-1)
    elif lr_scheduler_type == "polynominal":
        scheduler = get_polynomial_decay_schedule(optimizer, total_train_step, lr_end=config.lr_min, power=1.0, last_epoch=-1)
    elif lr_scheduler_type == "constant":
        scheduler = get_constant_schedule(optimizer)
    elif lr_scheduler_type == "constant_warmup":
        scheduler = get_constant_schedule_with_warmup(optimizer, config.warmup_steps)
    elif lr_scheduler_type == "linear_warmup":
        scheduler = get_linear_schedule_with_warmup(optimizer, config.warmup_steps, total_train_step)
    elif lr_scheduler_type == "one_cycle":
        step_per_epoch = int(total_train_step // config.num_train_epochs)
        scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=config.lr_max,
                                                        total_steps=total_train_step,
                                                        epochs=config.num_train_epochs,
                                                        steps_per_epoch=step_per_epoch,)
    elif lr_scheduler_type == "ladder":
        step_per_epoch = int(total_train_step // config.num_train_epochs)
        scheduler = get_ladder_schedule(optimizer, step_per_epoch, config.lr_min, )
    else:
        raise ValueError("Please  LR scheduler type. ")

    return scheduler


def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1):
    """
    Create a schedule with a constant learning rate, using the learning rate set in optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """
    return LambdaLR(optimizer, lambda _: 1, last_epoch=last_epoch)


def get_ladder_schedule(optimizer: Optimizer,
                        update_per_step: int,
                        lr_min: float,
                        decay_ratio: float = 0.95,
                        last_epoch: int = -1):
    """
    Create a schedule with a ladder learning rate.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    def lr_lambda(current_step: int):
        if current_step < update_per_step:
            return 1.0
        else:
            factor = math.floor(current_step / update_per_step)
            return float(decay_ratio ** factor)

    return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)


def get_constant_schedule_with_warmup(optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1):
    """
    Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate
    increases linearly between 0 and the initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1.0, num_warmup_steps))
        return 1.0

    return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)


def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):
    """
    Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0,
    after a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        return max(
            0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))
        )

    return LambdaLR(optimizer, lr_lambda, last_epoch)


def get_cosine_schedule_with_warmup(
    optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1
):
    """
    Create a schedule with a learning rate that decreases following the values of the cosine function between the
    initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the
    initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        num_cycles (:obj:`float`, `optional`, defaults to 0.5):
            The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0
            following a half-cosine).
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`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))
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))

    return LambdaLR(optimizer, lr_lambda, last_epoch)


def get_cosine_with_hard_restarts_schedule_with_warmup(
    optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1
):
    """
    Create a schedule with a learning rate that decreases following the values of the cosine function between the
    initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases
    linearly between 0 and the initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        num_cycles (:obj:`int`, `optional`, defaults to 1):
            The number of hard restarts to use.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`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))
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
        if progress >= 1.0:
            return 0.0
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0))))

    return LambdaLR(optimizer, lr_lambda, last_epoch)


def get_polynomial_decay_schedule_with_warmup(
    optimizer, num_warmup_steps, num_training_steps, lr_end=1e-7, power=1.0, last_epoch=-1
):
    """
    Create a schedule with a learning rate that decreases as a polynomial decay
    from the initial lr set in the optimizer to end lr defined by `lr_end`,
    after a warmup period during which it increases linearly from 0 to the
    initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        lr_end (:obj:`float`, `optional`, defaults to 1e-7):
            The end LR.
        power (:obj:`float`, `optional`, defaults to 1.0):
            Power factor.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Note: `power` defaults to 1.0 as in the fairseq implementation, which in turn is
    based on the original BERT implementation at
    https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    lr_init = optimizer.defaults["lr"]
    assert lr_init > lr_end, f"lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})"

    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        elif current_step > num_training_steps:
            return lr_end / lr_init  # as LambdaLR multiplies by lr_init
        else:
            lr_range = lr_init - lr_end
            decay_steps = num_training_steps - num_warmup_steps
            pct_remaining = 1 - (current_step - num_warmup_steps) / decay_steps
            decay = lr_range * pct_remaining ** power + lr_end
            return decay / lr_init  # as LambdaLR multiplies by lr_init

    return LambdaLR(optimizer, lr_lambda, last_epoch)



def get_polynomial_decay_schedule(
    optimizer, num_training_steps, lr_end=1e-7, power=1.0, last_epoch=-1
):
    """
    Create a schedule with a learning rate that decreases as a polynomial decay
    from the initial lr set in the optimizer to end lr defined by `lr_end`,
    after a warmup period during which it increases linearly from 0 to the
    initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        lr_end (:obj:`float`, `optional`, defaults to 1e-7):
            The end LR.
        power (:obj:`float`, `optional`, defaults to 1.0):
            Power factor.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Note: `power` defaults to 1.0 as in the fairseq implementation, which in turn is
    based on the original BERT implementation at
    https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    lr_init = optimizer.defaults["lr"]
    assert lr_init > lr_end, f"lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})"

    def lr_lambda(current_step: int):
        if current_step > num_training_steps:
            return lr_end / lr_init  # as LambdaLR multiplies by lr_init
        else:
            lr_range = lr_init - lr_end
            decay_steps = num_training_steps
            pct_remaining = 1 - (current_step ) / decay_steps
            decay = lr_range * pct_remaining ** power + lr_end
            return decay / lr_init  # as LambdaLR multiplies by lr_init

    return LambdaLR(optimizer, lr_lambda, last_epoch)
#!/usr/bin/env python3
# -*- coding: utf-8 -*-


# author: <NAME>
# file: lr_scheduler.py
# PyTorch optimization for BERT model.


import torch
import math
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR



def build_lr_scheduler(config, optimizer, total_train_step):
    lr_scheduler_type = config.lr_scheduler_type
    if lr_scheduler_type == "polynomial_warmup":
        scheduler = get_polynomial_decay_schedule_with_warmup(optimizer, config.warmup_steps, total_train_step,
                                                              lr_end=config.lr_min, power=1.0, last_epoch=-1)
    elif lr_scheduler_type == "polynominal":
        scheduler = get_polynomial_decay_schedule(optimizer, total_train_step, lr_end=config.lr_min, power=1.0, last_epoch=-1)
    elif lr_scheduler_type == "constant":
        scheduler = get_constant_schedule(optimizer)
    elif lr_scheduler_type == "constant_warmup":
        scheduler = get_constant_schedule_with_warmup(optimizer, config.warmup_steps)
    elif lr_scheduler_type == "linear_warmup":
        scheduler = get_linear_schedule_with_warmup(optimizer, config.warmup_steps, total_train_step)
    elif lr_scheduler_type == "one_cycle":
        step_per_epoch = int(total_train_step // config.num_train_epochs)
        scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=config.lr_max,
                                                        total_steps=total_train_step,
                                                        epochs=config.num_train_epochs,
                                                        steps_per_epoch=step_per_epoch,)
    elif lr_scheduler_type == "ladder":
        step_per_epoch = int(total_train_step // config.num_train_epochs)
        scheduler = get_ladder_schedule(optimizer, step_per_epoch, config.lr_min, )
    else:
        raise ValueError("Please  LR scheduler type. ")

    return scheduler


def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1):
    """
    Create a schedule with a constant learning rate, using the learning rate set in optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """
    return LambdaLR(optimizer, lambda _: 1, last_epoch=last_epoch)


def get_ladder_schedule(optimizer: Optimizer,
                        update_per_step: int,
                        lr_min: float,
                        decay_ratio: float = 0.95,
                        last_epoch: int = -1):
    """
    Create a schedule with a ladder learning rate.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    def lr_lambda(current_step: int):
        if current_step < update_per_step:
            return 1.0
        else:
            factor = math.floor(current_step / update_per_step)
            return float(decay_ratio ** factor)

    return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)


def get_constant_schedule_with_warmup(optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1):
    """
    Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate
    increases linearly between 0 and the initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1.0, num_warmup_steps))
        return 1.0

    return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)


def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):
    """
    Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0,
    after a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        return max(
            0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))
        )

    return LambdaLR(optimizer, lr_lambda, last_epoch)


def get_cosine_schedule_with_warmup(
    optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1
):
    """
    Create a schedule with a learning rate that decreases following the values of the cosine function between the
    initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the
    initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        num_cycles (:obj:`float`, `optional`, defaults to 0.5):
            The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0
            following a half-cosine).
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`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))
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))

    return LambdaLR(optimizer, lr_lambda, last_epoch)


def get_cosine_with_hard_restarts_schedule_with_warmup(
    optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1
):
    """
    Create a schedule with a learning rate that decreases following the values of the cosine function between the
    initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases
    linearly between 0 and the initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        num_cycles (:obj:`int`, `optional`, defaults to 1):
            The number of hard restarts to use.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Return:
        :obj:`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))
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
        if progress >= 1.0:
            return 0.0
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0))))

    return LambdaLR(optimizer, lr_lambda, last_epoch)


def get_polynomial_decay_schedule_with_warmup(
    optimizer, num_warmup_steps, num_training_steps, lr_end=1e-7, power=1.0, last_epoch=-1
):
    """
    Create a schedule with a learning rate that decreases as a polynomial decay
    from the initial lr set in the optimizer to end lr defined by `lr_end`,
    after a warmup period during which it increases linearly from 0 to the
    initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        lr_end (:obj:`float`, `optional`, defaults to 1e-7):
            The end LR.
        power (:obj:`float`, `optional`, defaults to 1.0):
            Power factor.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Note: `power` defaults to 1.0 as in the fairseq implementation, which in turn is
    based on the original BERT implementation at
    https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    lr_init = optimizer.defaults["lr"]
    assert lr_init > lr_end, f"lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})"

    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        elif current_step > num_training_steps:
            return lr_end / lr_init  # as LambdaLR multiplies by lr_init
        else:
            lr_range = lr_init - lr_end
            decay_steps = num_training_steps - num_warmup_steps
            pct_remaining = 1 - (current_step - num_warmup_steps) / decay_steps
            decay = lr_range * pct_remaining ** power + lr_end
            return decay / lr_init  # as LambdaLR multiplies by lr_init

    return LambdaLR(optimizer, lr_lambda, last_epoch)



def get_polynomial_decay_schedule(
    optimizer, num_training_steps, lr_end=1e-7, power=1.0, last_epoch=-1
):
    """
    Create a schedule with a learning rate that decreases as a polynomial decay
    from the initial lr set in the optimizer to end lr defined by `lr_end`,
    after a warmup period during which it increases linearly from 0 to the
    initial lr set in the optimizer.
    Args:
        optimizer (:class:`~torch.optim.Optimizer`):
            The optimizer for which to schedule the learning rate.
        num_warmup_steps (:obj:`int`):
            The number of steps for the warmup phase.
        num_training_steps (:obj:`int`):
            The total number of training steps.
        lr_end (:obj:`float`, `optional`, defaults to 1e-7):
            The end LR.
        power (:obj:`float`, `optional`, defaults to 1.0):
            Power factor.
        last_epoch (:obj:`int`, `optional`, defaults to -1):
            The index of the last epoch when resuming training.
    Note: `power` defaults to 1.0 as in the fairseq implementation, which in turn is
    based on the original BERT implementation at
    https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37
    Return:
        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
    """

    lr_init = optimizer.defaults["lr"]
    assert lr_init > lr_end, f"lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})"

    def lr_lambda(current_step: int):
        if current_step > num_training_steps:
            return lr_end / lr_init  # as LambdaLR multiplies by lr_init
        else:
            lr_range = lr_init - lr_end
            decay_steps = num_training_steps
            pct_remaining = 1 - (current_step ) / decay_steps
            decay = lr_range * pct_remaining ** power + lr_end
            return decay / lr_init  # as LambdaLR multiplies by lr_init

    return LambdaLR(optimizer, lr_lambda, last_epoch)
#!/usr/bin/env python3 # -*- coding: utf-8 -*- # author: <NAME> # file: lr_scheduler.py # PyTorch optimization for BERT model. Create a schedule with a constant learning rate, using the learning rate set in optimizer.     Args:         optimizer (:class:`~torch.optim.Optimizer`):             The optimizer for which to schedule the learning rate.         last_epoch (:obj:`int`, `optional`, defaults to -1):             The index of the last epoch when resuming training.     Return:         :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. Create a schedule with a ladder learning rate.     Args:         optimizer (:class:`~torch.optim.Optimizer`):             The optimizer for which to schedule the learning rate.         num_warmup_steps (:obj:`int`):             The number of steps for the warmup phase.         last_epoch (:obj:`int`, `optional`, defaults to -1):             The index of the last epoch when resuming training.     Return:         :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate     increases linearly between 0 and the initial lr set in the optimizer.     Args:         optimizer (:class:`~torch.optim.Optimizer`):             The optimizer for which to schedule the learning rate.         num_warmup_steps (:obj:`int`):             The number of steps for the warmup phase.         last_epoch (:obj:`int`, `optional`, defaults to -1):             The index of the last epoch when resuming training.     Return:         :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0,     after a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.     Args:         optimizer (:class:`~torch.optim.Optimizer`):             The optimizer for which to schedule the learning rate.         num_warmup_steps (:obj:`int`):             The number of steps for the warmup phase.         num_training_steps (:obj:`int`):             The total number of training steps.         last_epoch (:obj:`int`, `optional`, defaults to -1):             The index of the last epoch when resuming training.     Return:         :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. Create a schedule with a learning rate that decreases following the values of the cosine function between the     initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the     initial lr set in the optimizer.     Args:         optimizer (:class:`~torch.optim.Optimizer`):             The optimizer for which to schedule the learning rate.         num_warmup_steps (:obj:`int`):             The number of steps for the warmup phase.         num_training_steps (:obj:`int`):             The total number of training steps.         num_cycles (:obj:`float`, `optional`, defaults to 0.5):             The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0             following a half-cosine).         last_epoch (:obj:`int`, `optional`, defaults to -1):             The index of the last epoch when resuming training.     Return:         :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. Create a schedule with a learning rate that decreases following the values of the cosine function between the     initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases     linearly between 0 and the initial lr set in the optimizer.     Args:         optimizer (:class:`~torch.optim.Optimizer`):             The optimizer for which to schedule the learning rate.         num_warmup_steps (:obj:`int`):             The number of steps for the warmup phase.         num_training_steps (:obj:`int`):             The total number of training steps.         num_cycles (:obj:`int`, `optional`, defaults to 1):             The number of hard restarts to use.         last_epoch (:obj:`int`, `optional`, defaults to -1):             The index of the last epoch when resuming training.     Return:         :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. Create a schedule with a learning rate that decreases as a polynomial decay     from the initial lr set in the optimizer to end lr defined by `lr_end`,     after a warmup period during which it increases linearly from 0 to the     initial lr set in the optimizer.     Args:         optimizer (:class:`~torch.optim.Optimizer`):             The optimizer for which to schedule the learning rate.         num_warmup_steps (:obj:`int`):             The number of steps for the warmup phase.         num_training_steps (:obj:`int`):             The total number of training steps.         lr_end (:obj:`float`, `optional`, defaults to 1e-7):             The end LR.         power (:obj:`float`, `optional`, defaults to 1.0):             Power factor.         last_epoch (:obj:`int`, `optional`, defaults to -1):             The index of the last epoch when resuming training.     Note: `power` defaults to 1.0 as in the fairseq implementation, which in turn is     based on the original BERT implementation at     https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37     Return:         :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. # as LambdaLR multiplies by lr_init # as LambdaLR multiplies by lr_init Create a schedule with a learning rate that decreases as a polynomial decay     from the initial lr set in the optimizer to end lr defined by `lr_end`,     after a warmup period during which it increases linearly from 0 to the     initial lr set in the optimizer.     Args:         optimizer (:class:`~torch.optim.Optimizer`):             The optimizer for which to schedule the learning rate.         num_warmup_steps (:obj:`int`):             The number of steps for the warmup phase.         num_training_steps (:obj:`int`):             The total number of training steps.         lr_end (:obj:`float`, `optional`, defaults to 1e-7):             The end LR.         power (:obj:`float`, `optional`, defaults to 1.0):             Power factor.         last_epoch (:obj:`int`, `optional`, defaults to -1):             The index of the last epoch when resuming training.     Note: `power` defaults to 1.0 as in the fairseq implementation, which in turn is     based on the original BERT implementation at     https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37     Return:         :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. # as LambdaLR multiplies by lr_init # as LambdaLR multiplies by lr_init