text: # Code referenced from https://github.com/Bjarten/early-stopping-pytorch/blob/master/pytorchtools.py
  
  import numpy as np
  import torch
  import math
  from torch.optim.optimizer import Optimizer
  from torch.optim.lr_scheduler import _LRScheduler
  
  
  class EarlyStopping:
      """Early stops the training if validation loss doesn't improve after a given patience."""
      def __init__(self, save_name, patience=7, verbose=False, delta=0):
          """
          Args:
              patience (int): How long to wait after last time validation loss improved.
                              Default: 7
              verbose (bool): If True, prints a message for each validation loss improvement.
                              Default: False
              delta (float): Minimum change in the monitored quantity to qualify as an improvement.
                              Default: 0
          """
          self.patience = patience
          self.verbose = verbose
          self.counter = 0
          self.best_score = None
          self.early_stop = False
          self.val_loss_min = np.Inf
          self.delta = delta
          self.save_name = save_name
  
      def __call__(self, val_loss, model):
  
          score = -val_loss
  
          if self.best_score is None:
              self.best_score = score
              self.save_checkpoint(val_loss, model)
          elif score < self.best_score - self.delta:
              self.counter += 1
              print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
              if self.counter >= self.patience:
                  self.early_stop = True
          else:
              self.best_score = score
              self.save_checkpoint(val_loss, model)
              self.counter = 0
  
      def save_checkpoint(self, val_loss, model):
          '''Saves model when validation loss decrease.'''
          if self.verbose:
              print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}).  Saving model...')
          torch.save(model.state_dict(), 'output/weights/' + self.save_name + '.pth')
          self.val_loss_min = val_loss
  
  
  # from https://github.com/mpyrozhok/adamwr
  
  
  class AdamW(Optimizer):
      """Implements Adam algorithm.
      Arguments:
          params (iterable): iterable of parameters to optimize or dicts defining
              parameter groups
          lr (float, optional): learning rate (default: 1e-3)
          betas (Tuple[float, float], optional): coefficients used for computing
              running averages of gradient and its square (default: (0.9, 0.999))
          eps (float, optional): term added to the denominator to improve
              numerical stability (default: 1e-8)
          weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
          amsgrad (boolean, optional): whether to use the AMSGrad variant of this
              algorithm from the paper `On the Convergence of Adam and Beyond`_
      """
  
      def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                   weight_decay=0, amsgrad=False):
          if not 0.0 <= betas[0] < 1.0:
              raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
          if not 0.0 <= betas[1] < 1.0:
              raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
          defaults = dict(lr=lr, betas=betas, eps=eps,
                          weight_decay=weight_decay, amsgrad=amsgrad)
          #super(AdamW, self).__init__(params, defaults)
          super().__init__(params, defaults)
  
      def step(self, closure=None):
          """Performs a single optimization step.
          Arguments:
              closure (callable, optional): A closure that reevaluates the model
                  and returns the loss.
          """
          loss = None
          if closure is not None:
              loss = closure()
  
          for group in self.param_groups:
              for p in group['params']:
                  if p.grad is None:
                      continue
                  grad = p.grad.data
                  if grad.is_sparse:
                      raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
                  amsgrad = group['amsgrad']
  
                  state = self.state[p]
  
                  # State initialization
                  if len(state) == 0:
                      state['step'] = 0
                      # Exponential moving average of gradient values
                      state['exp_avg'] = torch.zeros_like(p.data)
                      # Exponential moving average of squared gradient values
                      state['exp_avg_sq'] = torch.zeros_like(p.data)
                      if amsgrad:
                          # Maintains max of all exp. moving avg. of sq. grad. values
                          state['max_exp_avg_sq'] = torch.zeros_like(p.data)
  
                  exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                  if amsgrad:
                      max_exp_avg_sq = state['max_exp_avg_sq']
                  beta1, beta2 = group['betas']
  
                  state['step'] += 1
  
                  # Decay the first and second moment running average coefficient
                  exp_avg.mul_(beta1).add_(1 - beta1, grad)
                  exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
                  if amsgrad:
                      # Maintains the maximum of all 2nd moment running avg. till now
                      torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
                      # Use the max. for normalizing running avg. of gradient
                      denom = max_exp_avg_sq.sqrt().add_(group['eps'])
                  else:
                      denom = exp_avg_sq.sqrt().add_(group['eps'])
  
                  bias_correction1 = 1 - beta1 ** state['step']
                  bias_correction2 = 1 - beta2 ** state['step']
                  step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1
  
                  p.data.mul_(1 - group['weight_decay']).addcdiv_(-step_size, exp_avg, denom)
  
          return loss
  
  
  
  class ReduceMaxLROnRestart:
      def __init__(self, ratio=0.75):
          self.ratio = ratio
  
      def __call__(self, eta_min, eta_max):
          return eta_min, eta_max * self.ratio
  
  
  class ExpReduceMaxLROnIteration:
      def __init__(self, gamma=1):
          self.gamma = gamma
  
      def __call__(self, eta_min, eta_max, iterations):
          return eta_min, eta_max * self.gamma ** iterations
  
  
  class CosinePolicy:
      def __call__(self, t_cur, restart_period):
          return 0.5 * (1. + math.cos(math.pi *
                                      (t_cur / restart_period)))
  
  
  class ArccosinePolicy:
      def __call__(self, t_cur, restart_period):
          return (math.acos(max(-1, min(1, 2 * t_cur
                                        / restart_period - 1))) / math.pi)
  
  
  class TriangularPolicy:
      def __init__(self, triangular_step=0.5):
          self.triangular_step = triangular_step
  
      def __call__(self, t_cur, restart_period):
          inflection_point = self.triangular_step * restart_period
          point_of_triangle = (t_cur / inflection_point
                               if t_cur < inflection_point
                               else 1.0 - (t_cur - inflection_point)
                               / (restart_period - inflection_point))
          return point_of_triangle
  
  
  class CyclicLRWithRestarts(_LRScheduler):
      """Decays learning rate with cosine annealing, normalizes weight decay
      hyperparameter value, implements restarts.
      https://arxiv.org/abs/1711.05101
      Args:
          optimizer (Optimizer): Wrapped optimizer.
          batch_size: minibatch size
          epoch_size: training samples per epoch
          restart_period: epoch count in the first restart period
          t_mult: multiplication factor by which the next restart period will expand/shrink
          policy: ["cosine", "arccosine", "triangular", "triangular2", "exp_range"]
          min_lr: minimum allowed learning rate
          verbose: print a message on every restart
          gamma: exponent used in "exp_range" policy
          eta_on_restart_cb: callback executed on every restart, adjusts max or min lr
          eta_on_iteration_cb: callback executed on every iteration, adjusts max or min lr
          triangular_step: adjusts ratio of increasing/decreasing phases for triangular policy
      """
  
      def __init__(self, optimizer, batch_size, epoch_size, restart_period=100,
                   t_mult=2, last_epoch=-1, verbose=False,
                   policy="cosine", policy_fn=None, min_lr=1e-7,
                   eta_on_restart_cb=None, eta_on_iteration_cb=None,
                   gamma=1.0, triangular_step=0.5):
  
          if not isinstance(optimizer, Optimizer):
              raise TypeError('{} is not an Optimizer'.format(
                  type(optimizer).__name__))
  
          self.optimizer = optimizer
  
          if last_epoch == -1:
              for group in optimizer.param_groups:
                  group.setdefault('initial_lr', group['lr'])
                  group.setdefault('minimum_lr', min_lr)
          else:
              for i, group in enumerate(optimizer.param_groups):
                  if 'initial_lr' not in group:
                      raise KeyError("param 'initial_lr' is not specified "
                                     "in param_groups[{}] when resuming an"
                                     " optimizer".format(i))
  
          self.base_lrs = [group['initial_lr'] for group
                           in optimizer.param_groups]
  
          self.min_lrs = [group['minimum_lr'] for group
                          in optimizer.param_groups]
  
          self.base_weight_decays = [group['weight_decay'] for group
                                     in optimizer.param_groups]
  
          self.policy = policy
          self.eta_on_restart_cb = eta_on_restart_cb
          self.eta_on_iteration_cb = eta_on_iteration_cb
          if policy_fn is not None:
              self.policy_fn = policy_fn
          elif self.policy == "cosine":
              self.policy_fn = CosinePolicy()
          elif self.policy == "arccosine":
              self.policy_fn = ArccosinePolicy()
          elif self.policy == "triangular":
              self.policy_fn = TriangularPolicy(triangular_step=triangular_step)
          elif self.policy == "triangular2":
              self.policy_fn = TriangularPolicy(triangular_step=triangular_step)
              self.eta_on_restart_cb = ReduceMaxLROnRestart(ratio=0.5)
          elif self.policy == "exp_range":
              self.policy_fn = TriangularPolicy(triangular_step=triangular_step)
              self.eta_on_iteration_cb = ExpReduceMaxLROnIteration(gamma=gamma)
  
          self.last_epoch = last_epoch
          self.batch_size = batch_size
          self.epoch_size = epoch_size
  
          self.iteration = 0
          self.total_iterations = 0
  
          self.t_mult = t_mult
          self.verbose = verbose
          self.restart_period = math.ceil(restart_period)
          self.restarts = 0
          self.t_epoch = -1
          self.epoch = -1
  
          self.eta_min = 0
          self.eta_max = 1
  
          self.end_of_period = False
          self.batch_increments = []
          self._set_batch_increment()
  
          self.t_cur = None
  
      def _on_restart(self):
          if self.eta_on_restart_cb is not None:
              self.eta_min, self.eta_max = self.eta_on_restart_cb(self.eta_min,
                                                                  self.eta_max)
  
      def _on_iteration(self):
          if self.eta_on_iteration_cb is not None:
              self.eta_min, self.eta_max = self.eta_on_iteration_cb(self.eta_min,
                                                                    self.eta_max,
                                                                    self.total_iterations)
  
      def get_lr(self, t_cur):
          eta_t = (self.eta_min + (self.eta_max - self.eta_min)
                   * self.policy_fn(t_cur, self.restart_period))
  
          weight_decay_norm_multi = math.sqrt(self.batch_size /
                                              (self.epoch_size *
                                               self.restart_period))
  
          lrs = [min_lr + (base_lr - min_lr) * eta_t for base_lr, min_lr
                 in zip(self.base_lrs, self.min_lrs)]
          weight_decays = [base_weight_decay * eta_t * weight_decay_norm_multi
                           for base_weight_decay in self.base_weight_decays]
  
          if (self.t_epoch + 1) % self.restart_period < self.t_epoch:
              self.end_of_period = True
  
          if self.t_epoch % self.restart_period < self.t_epoch:
              if self.verbose:
                  print("Restart {} at epoch {}".format(self.restarts + 1,
                                                        self.last_epoch))
              self.restart_period = math.ceil(self.restart_period * self.t_mult)
              self.restarts += 1
              self.t_epoch = 0
              self._on_restart()
              self.end_of_period = False
  
          return zip(lrs, weight_decays)
  
      def _set_batch_increment(self):
          d, r = divmod(self.epoch_size, self.batch_size)
          batches_in_epoch = d + 2 if r > 0 else d + 1
          self.iteration = 0
          self.batch_increments = torch.linspace(0, 1, batches_in_epoch).tolist()
  
      def step(self):
          self.last_epoch += 1
          self.t_epoch += 1
          self._set_batch_increment()
          self.batch_step()
  
      def batch_step(self):
          try:
              t_cur = self.t_epoch + self.batch_increments[self.iteration]
              self.t_cur = t_cur
              self._on_iteration()
              self.iteration += 1
              self.total_iterations += 1
          except IndexError:
              raise StopIteration("Epoch size and batch size used in the "
                                  "training loop and while initializing "
                                  "scheduler should be the same.")
  
          for param_group, (lr, weight_decay) in zip(self.optimizer.param_groups,
                                                     self.get_lr(t_cur)):
              param_group['lr'] = lr
              param_group['weight_decay'] = weight_decay
  