content: import math
  import warnings
  from collections import Counter
  from bisect import bisect_right
  
  import torch
  from torch.optim.lr_scheduler import _LRScheduler
  
  
  class LinearWarmupCosineAnnealingLR(_LRScheduler):
      """
      Sets the learning rate of each parameter group to follow a linear warmup schedule
      between warmup_start_lr and base_lr followed by a cosine annealing schedule between
      base_lr and eta_min.
  
      .. warning::
          It is recommended to call :func:`.step()` for :class:`LinearWarmupCosineAnnealingLR`
          after each iteration as calling it after each epoch will keep the starting lr at
          warmup_start_lr for the first epoch which is 0 in most cases.
  
      .. warning::
          passing epoch to :func:`.step()` is being deprecated and comes with an EPOCH_DEPRECATION_WARNING.
          It calls the :func:`_get_closed_form_lr()` method for this scheduler instead of
          :func:`get_lr()`. Though this does not change the behavior of the scheduler, when passing
          epoch param to :func:`.step()`, the user should call the :func:`.step()` function before calling
          train and validation methods.
  
      Example:
          >>> layer = nn.Linear(10, 1)
          >>> optimizer = Adam(layer.parameters(), lr=0.02)
          >>> scheduler = LinearWarmupCosineAnnealingLR(optimizer, warmup_epochs=10, max_epochs=40)
          >>> #
          >>> # the default case
          >>> for epoch in range(40):
          ...     # train(...)
          ...     # validate(...)
          ...     scheduler.step()
          >>> #
          >>> # passing epoch param case
          >>> for epoch in range(40):
          ...     scheduler.step(epoch)
          ...     # train(...)
          ...     # validate(...)
      """
  
      def __init__(
          self,
          optimizer,
          warmup_epochs,
          max_epochs,
          warmup_start_lr = 0.0,
          eta_min = 1e-8,
          last_epoch = -1,
      ):
          """
          Args:
              optimizer (Optimizer): Wrapped optimizer.
              warmup_epochs (int): Maximum number of iterations for linear warmup
              max_epochs (int): Maximum number of iterations
              warmup_start_lr (float): Learning rate to start the linear warmup. Default: 0.
              eta_min (float): Minimum learning rate. Default: 0.
              last_epoch (int): The index of last epoch. Default: -1.
          """
          self.warmup_epochs = warmup_epochs
          self.max_epochs = max_epochs
          self.warmup_start_lr = warmup_start_lr
          self.eta_min = eta_min
  
          super(LinearWarmupCosineAnnealingLR, self).__init__(optimizer, last_epoch)
  
      def get_lr(self):
          """
          Compute learning rate using chainable form of the scheduler
          """
          if not self._get_lr_called_within_step:
              warnings.warn(
                  "To get the last learning rate computed by the scheduler, "
                  "please use `get_last_lr()`.",
                  UserWarning,
              )
  
          if self.last_epoch == 0:
              return [self.warmup_start_lr] * len(self.base_lrs)
          elif self.last_epoch < self.warmup_epochs:
              return [
                  group["lr"] + (base_lr - self.warmup_start_lr) / (self.warmup_epochs - 1)
                  for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
              ]
          elif self.last_epoch == self.warmup_epochs:
              return self.base_lrs
          elif (self.last_epoch - 1 - self.max_epochs) % (2 * (self.max_epochs - self.warmup_epochs)) == 0:
              return [
                  group["lr"] + (base_lr - self.eta_min) *
                  (1 - math.cos(math.pi / (self.max_epochs - self.warmup_epochs))) / 2
                  for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
              ]
  
          return [
              (1 + math.cos(math.pi * (self.last_epoch - self.warmup_epochs) / (self.max_epochs - self.warmup_epochs))) /
              (
                  1 +
                  math.cos(math.pi * (self.last_epoch - self.warmup_epochs - 1) / (self.max_epochs - self.warmup_epochs))
              ) * (group["lr"] - self.eta_min) + self.eta_min for group in self.optimizer.param_groups
          ]
  
      def _get_closed_form_lr(self):
          """
          Called when epoch is passed as a param to the `step` function of the scheduler.
          """
          if self.last_epoch < self.warmup_epochs:
              return [
                  self.warmup_start_lr + self.last_epoch * (base_lr - self.warmup_start_lr) / (self.warmup_epochs - 1)
                  for base_lr in self.base_lrs
              ]
  
          return [
              self.eta_min + 0.5 * (base_lr - self.eta_min) *
              (1 + math.cos(math.pi * (self.last_epoch - self.warmup_epochs) / (self.max_epochs - self.warmup_epochs)))
              for base_lr in self.base_lrs
          ]
  
  
  class LinearWarmupMultiStepLR(_LRScheduler):
      """
      Sets the learning rate of each parameter group to follow a linear warmup schedule
      between warmup_start_lr and base_lr followed by a multi-step schedule that decays
      the learning rate of each parameter group by gamma once the
      number of epoch reaches one of the milestones.
  
      .. warning::
          It is recommended to call :func:`.step()` for :class:`LinearWarmupCosineAnnealingLR`
          after each iteration as calling it after each epoch will keep the starting lr at
          warmup_start_lr for the first epoch which is 0 in most cases.
  
      .. warning::
          passing epoch to :func:`.step()` is being deprecated and comes with an EPOCH_DEPRECATION_WARNING.
          It calls the :func:`_get_closed_form_lr()` method for this scheduler instead of
          :func:`get_lr()`. Though this does not change the behavior of the scheduler, when passing
          epoch param to :func:`.step()`, the user should call the :func:`.step()` function before calling
          train and validation methods.
      """
  
      def __init__(
          self,
          optimizer,
          warmup_epochs,
          milestones,
          warmup_start_lr = 0.0,
          gamma = 0.1,
          last_epoch = -1,
      ):
          """
          Args:
              optimizer (Optimizer): Wrapped optimizer.
              warmup_epochs (int): Maximum number of iterations for linear warmup
              max_epochs (int): Maximum number of iterations
              milestones (list): List of epoch indices. Must be increasing.
              warmup_start_lr (float): Learning rate to start the linear warmup. Default: 0.
              gamma (float): Multiplicative factor of learning rate decay.
              Default: 0.1.
              last_epoch (int): The index of last epoch. Default: -1.
          """
          self.warmup_epochs = warmup_epochs
          self.warmup_start_lr = warmup_start_lr
          self.milestones = Counter(milestones)
          self.gamma = gamma
  
          super(LinearWarmupMultiStepLR, self).__init__(optimizer, last_epoch)
  
      def get_lr(self):
          """
          Compute learning rate using chainable form of the scheduler
          """
          if not self._get_lr_called_within_step:
              warnings.warn("To get the last learning rate computed by the scheduler, "
                            "please use `get_last_lr()`.", UserWarning)
  
          if self.last_epoch == 0:
              # starting warm up
              return [self.warmup_start_lr] * len(self.base_lrs)
          elif self.last_epoch < self.warmup_epochs:
              # linear warm up (0 ~ self.warmup_epochs -1)
              return [
                  group["lr"] + (base_lr - self.warmup_start_lr) / (self.warmup_epochs - 1)
                  for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
              ]
          elif self.last_epoch == self.warmup_epochs:
              # end of warm up (reset to base lrs)
              return self.base_lrs
          elif (self.last_epoch - self.warmup_epochs) not in self.milestones:
              # in between the steps
              return [group['lr'] for group in self.optimizer.param_groups]
  
          return [
              group['lr'] * self.gamma ** self.milestones[self.last_epoch - self.warmup_epochs]
              for group in self.optimizer.param_groups
          ]
  
      def _get_closed_form_lr(self):
          """
          Called when epoch is passed as a param to the `step` function of the scheduler.
          """
          if self.last_epoch < self.warmup_epochs:
              return [
                  self.warmup_start_lr + self.last_epoch * (base_lr - self.warmup_start_lr) / (self.warmup_epochs - 1)
                  for base_lr in self.base_lrs
              ]
  
          milestones = list(sorted(self.milestones.elements()))
          return [base_lr * self.gamma ** bisect_right(milestones, self.last_epoch - self.warmup_epochs)
                  for base_lr in self.base_lrs]
  