text: # 
  """
  Reference: https://github.com/davidtvs/pytorch-lr-finder
  """
  
  from __future__ import print_function, with_statement, division
  import copy
  import os
  import torch
  from tqdm.autonotebook import tqdm
  from torch.optim.lr_scheduler import _LRScheduler
  import matplotlib.pyplot as plt
  
  try:
      from apex import amp
  
      IS_AMP_AVAILABLE = True
  except ImportError:
      import logging
  
      logging.basicConfig()
      logger = logging.getLogger(__name__)
      logger.warning(
          "To enable mixed precision training, please install `apex`. "
          "Or you can re-install this package by the following command:\n"
          '  pip install torch-lr-finder -v --global-option="amp"'
      )
      IS_AMP_AVAILABLE = False
      del logging
  
  
  class LRFinder(object):
      """Learning rate range test.
      The learning rate range test increases the learning rate in a pre-training run
      between two boundaries in a linear or exponential manner. It provides valuable
      information on how well the network can be trained over a range of learning rates
      and what is the optimal learning rate.
      Arguments:
          model (torch.nn.Module): wrapped model.
          optimizer (torch.optim.Optimizer): wrapped optimizer where the defined learning
              is assumed to be the lower boundary of the range test.
          criterion (torch.nn.Module): wrapped loss function.
          device (str or torch.device, optional): a string ("cpu" or "cuda") with an
              optional ordinal for the device type (e.g. "cuda:X", where is the ordinal).
              Alternatively, can be an object representing the device on which the
              computation will take place. Default: None, uses the same device as `model`.
          memory_cache (boolean, optional): if this flag is set to True, `state_dict` of
              model and optimizer will be cached in memory. Otherwise, they will be saved
              to files under the `cache_dir`.
          cache_dir (string, optional): path for storing temporary files. If no path is
              specified, system-wide temporary directory is used. Notice that this
              parameter will be ignored if `memory_cache` is True.
      Example:
          >>> lr_finder = LRFinder(net, optimizer, criterion, device="cuda")
          >>> lr_finder.range_test(dataloader, end_lr=100, num_iter=100)
          >>> lr_finder.plot() # to inspect the loss-learning rate graph
          >>> lr_finder.reset() # to reset the model and optimizer to their initial state
      Reference:
      Cyclical Learning Rates for Training Neural Networks: https://arxiv.org/abs/1506.01186
      fastai/lr_find: https://github.com/fastai/fastai
      """
  
      def __init__(
          self,
          model,
          optimizer,
          criterion,
          device=None,
          memory_cache=True,
          cache_dir=None,
      ):
          # Check if the optimizer is already attached to a scheduler
          self.optimizer = optimizer
          self._check_for_scheduler()
  
          self.model = model
          self.criterion = criterion
          self.history = {"lr": [], "loss": []}
          self.best_loss = None
          self.memory_cache = memory_cache
          self.cache_dir = cache_dir
  
          # Save the original state of the model and optimizer so they can be restored if
          # needed
          self.model_device = next(self.model.parameters()).device
          self.state_cacher = StateCacher(memory_cache, cache_dir=cache_dir)
          self.state_cacher.store("model", self.model.state_dict())
          self.state_cacher.store("optimizer", self.optimizer.state_dict())
  
          # If device is None, use the same as the model
          if device:
              self.device = device
          else:
              self.device = self.model_device
  
      def reset(self):
          """Restores the model and optimizer to their initial states."""
  
          self.model.load_state_dict(self.state_cacher.retrieve("model"))
          self.optimizer.load_state_dict(self.state_cacher.retrieve("optimizer"))
          self.model.to(self.model_device)
  
      def range_test(
          self,
          train_loader,
          val_loader=None,
          start_lr=None,
          end_lr=10,
          num_iter=100,
          step_mode="exp",
          smooth_f=0.05,
          diverge_th=5,
          accumulation_steps=1,
      ):
          """Performs the learning rate range test.
          Arguments:
              train_loader (torch.utils.data.DataLoader): the training set data laoder.
              val_loader (torch.utils.data.DataLoader, optional): if `None` the range test
                  will only use the training loss. When given a data loader, the model is
                  evaluated after each iteration on that dataset and the evaluation loss
                  is used. Note that in this mode the test takes significantly longer but
                  generally produces more precise results. Default: None.
              start_lr (float, optional): the starting learning rate for the range test.
                  Default: None (uses the learning rate from the optimizer).
              end_lr (float, optional): the maximum learning rate to test. Default: 10.
              num_iter (int, optional): the number of iterations over which the test
                  occurs. Default: 100.
              step_mode (str, optional): one of the available learning rate policies,
                  linear or exponential ("linear", "exp"). Default: "exp".
              smooth_f (float, optional): the loss smoothing factor within the [0, 1[
                  interval. Disabled if set to 0, otherwise the loss is smoothed using
                  exponential smoothing. Default: 0.05.
              diverge_th (int, optional): the test is stopped when the loss surpasses the
                  threshold:  diverge_th * best_loss. Default: 5.
              accumulation_steps (int, optional): steps for gradient accumulation. If it
                  is 1, gradients are not accumulated. Default: 1.
          Example (fastai approach):
              >>> lr_finder = LRFinder(net, optimizer, criterion, device="cuda")
              >>> lr_finder.range_test(dataloader, end_lr=100, num_iter=100)
          Example (Leslie Smith's approach):
              >>> lr_finder = LRFinder(net, optimizer, criterion, device="cuda")
              >>> lr_finder.range_test(trainloader, val_loader=val_loader, end_lr=1, num_iter=100, step_mode="linear")
          Gradient accumulation is supported; example:
              >>> train_data = ...    # prepared dataset
              >>> desired_bs, real_bs = 32, 4         # batch size
              >>> accumulation_steps = desired_bs // real_bs     # required steps for accumulation
              >>> dataloader = torch.utils.data.DataLoader(train_data, batch_size=real_bs, shuffle=True)
              >>> acc_lr_finder = LRFinder(net, optimizer, criterion, device="cuda")
              >>> acc_lr_finder.range_test(dataloader, end_lr=10, num_iter=100, accumulation_steps=accumulation_steps)
          Reference:
          [Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups](
          https://medium.com/huggingface/ec88c3e51255)
          [thomwolf/gradient_accumulation](https://gist.github.com/thomwolf/ac7a7da6b1888c2eeac8ac8b9b05d3d3)
          """
  
          # Reset test results
          self.history = {"lr": [], "loss": []}
          self.best_loss = None
  
          # Move the model to the proper device
          self.model.to(self.device)
  
          # Check if the optimizer is already attached to a scheduler
          self._check_for_scheduler()
  
          # Set the starting learning rate
          if start_lr:
              self._set_learning_rate(start_lr)
  
          # Initialize the proper learning rate policy
          if step_mode.lower() == "exp":
              lr_schedule = ExponentialLR(self.optimizer, end_lr, num_iter)
          elif step_mode.lower() == "linear":
              lr_schedule = LinearLR(self.optimizer, end_lr, num_iter)
          else:
              raise ValueError("expected one of (exp, linear), got {}".format(step_mode))
  
          if smooth_f < 0 or smooth_f >= 1:
              raise ValueError("smooth_f is outside the range [0, 1[")
  
          # Create an iterator to get data batch by batch
          iter_wrapper = DataLoaderIterWrapper(train_loader)
          for iteration in tqdm(range(num_iter)):
              # Train on batch and retrieve loss
              loss = self._train_batch(iter_wrapper, accumulation_steps)
              if val_loader:
                  loss = self._validate(val_loader)
  
              # Update the learning rate
              lr_schedule.step()
              self.history["lr"].append(lr_schedule.get_lr()[0])
  
              # Track the best loss and smooth it if smooth_f is specified
              if iteration == 0:
                  self.best_loss = loss
              else:
                  if smooth_f > 0:
                      loss = smooth_f * loss + (1 - smooth_f) * self.history["loss"][-1]
                  if loss < self.best_loss:
                      self.best_loss = loss
  
              # Check if the loss has diverged; if it has, stop the test
              self.history["loss"].append(loss)
              if loss > diverge_th * self.best_loss:
                  print("Stopping early, the loss has diverged")
                  break
  
          print("Learning rate search finished. See the graph with {finder_name}.plot()")
  
      def _set_learning_rate(self, new_lrs):
          if not isinstance(new_lrs, list):
              new_lrs = [new_lrs] * len(self.optimizer.param_groups)
          if len(new_lrs) != len(self.optimizer.param_groups):
              raise ValueError(
                  "Length of `new_lrs` is not equal to the number of parameter groups "
                  + "in the given optimizer"
              )
  
          for param_group, new_lr in zip(self.optimizer.param_groups, new_lrs):
              param_group["lr"] = new_lr
  
      def _check_for_scheduler(self):
          for param_group in self.optimizer.param_groups:
              if "initial_lr" in param_group:
                  raise RuntimeError("Optimizer already has a scheduler attached to it")
  
      def _train_batch(self, iter_wrapper, accumulation_steps):
          self.model.train()
          total_loss = None  # for late initialization
  
          self.optimizer.zero_grad()
          for i in range(accumulation_steps):
              inputs, labels = iter_wrapper.get_batch()
              inputs, labels = self._move_to_device(inputs, labels)
  
              # Forward pass
              outputs = self.model(inputs)
              loss = self.criterion(outputs, labels)
  
              # Loss should be averaged in each step
              loss /= accumulation_steps
  
              # Backward pass
              if IS_AMP_AVAILABLE and hasattr(self.optimizer, "_amp_stash"):
                  # For minor performance optimization, see also:
                  # https://nvidia.github.io/apex/advanced.html#gradient-accumulation-across-iterations
                  delay_unscale = ((i + 1) % accumulation_steps) != 0
  
                  with amp.scale_loss(
                      loss, self.optimizer, delay_unscale=delay_unscale
                  ) as scaled_loss:
                      scaled_loss.backward()
              else:
                  loss.backward()
  
              if total_loss is None:
                  total_loss = loss
              else:
                  total_loss += loss
  
          self.optimizer.step()
  
          return total_loss.item()
  
      def _move_to_device(self, inputs, labels):
          def move(obj, device):
              if isinstance(obj, tuple):
                  return tuple(move(o, device) for o in obj)
              elif torch.is_tensor(obj):
                  return obj.to(device)
              elif isinstance(obj, list):
                  return [move(o, device) for o in obj]
              else:
                  return obj
  
          inputs = move(inputs, self.device)
          labels = move(labels, self.device)
          return inputs, labels
  
      def _validate(self, dataloader):
          # Set model to evaluation mode and disable gradient computation
          running_loss = 0
          self.model.eval()
          with torch.no_grad():
              for inputs, labels in dataloader:
                  # Move data to the correct device
                  inputs, labels = self._move_to_device(inputs, labels)
  
                  if isinstance(inputs, tuple) or isinstance(inputs, list):
                      batch_size = inputs[0].size(0)
                  else:
                      batch_size = inputs.size(0)
  
                  # Forward pass and loss computation
                  outputs = self.model(inputs)
                  loss = self.criterion(outputs, labels)
                  running_loss += loss.item() * batch_size
  
          return running_loss / len(dataloader.dataset)
  
      def plot(self, skip_start=10, skip_end=5, log_lr=True, show_lr=None, ax=None):
          """Plots the learning rate range test.
          Arguments:
              skip_start (int, optional): number of batches to trim from the start.
                  Default: 10.
              skip_end (int, optional): number of batches to trim from the start.
                  Default: 5.
              log_lr (bool, optional): True to plot the learning rate in a logarithmic
                  scale; otherwise, plotted in a linear scale. Default: True.
              show_lr (float, optional): if set, adds a vertical line to visualize the
                  specified learning rate. Default: None.
              ax (matplotlib.axes.Axes, optional): the plot is created in the specified
                  matplotlib axes object and the figure is not be shown. If `None`, then
                  the figure and axes object are created in this method and the figure is
                  shown . Default: None.
          Returns:
              The matplotlib.axes.Axes object that contains the plot.
          """
  
          if skip_start < 0:
              raise ValueError("skip_start cannot be negative")
          if skip_end < 0:
              raise ValueError("skip_end cannot be negative")
          if show_lr is not None and not isinstance(show_lr, float):
              raise ValueError("show_lr must be float")
  
          # Get the data to plot from the history dictionary. Also, handle skip_end=0
          # properly so the behaviour is the expected
          lrs = self.history["lr"]
          losses = self.history["loss"]
          if skip_end == 0:
              lrs = lrs[skip_start:]
              losses = losses[skip_start:]
          else:
              lrs = lrs[skip_start:-skip_end]
              losses = losses[skip_start:-skip_end]
  
          # Create the figure and axes object if axes was not already given
          fig = None
          if ax is None:
              fig, ax = plt.subplots()
  
          # Plot loss as a function of the learning rate
          ax.plot(lrs, losses)
          if log_lr:
              ax.set_xscale("log")
          ax.set_xlabel("Learning rate")
          ax.set_ylabel("Loss")
  
          if show_lr is not None:
              ax.axvline(x=show_lr, color="red")
  
          # Show only if the figure was created internally
          if fig is not None:
              plt.show()
  
          return ax
  
  
  class LinearLR(_LRScheduler):
      """Linearly increases the learning rate between two boundaries over a number of
      iterations.
      Arguments:
          optimizer (torch.optim.Optimizer): wrapped optimizer.
          end_lr (float): the final learning rate.
          num_iter (int): the number of iterations over which the test occurs.
          last_epoch (int, optional): the index of last epoch. Default: -1.
      """
  
      def __init__(self, optimizer, end_lr, num_iter, last_epoch=-1):
          self.end_lr = end_lr
          self.num_iter = num_iter
          super(LinearLR, self).__init__(optimizer, last_epoch)
  
      def get_lr(self):
          curr_iter = self.last_epoch + 1
          r = curr_iter / self.num_iter
          return [base_lr + r * (self.end_lr - base_lr) for base_lr in self.base_lrs]
  
  
  class ExponentialLR(_LRScheduler):
      """Exponentially increases the learning rate between two boundaries over a number of
      iterations.
      Arguments:
          optimizer (torch.optim.Optimizer): wrapped optimizer.
          end_lr (float): the final learning rate.
          num_iter (int): the number of iterations over which the test occurs.
          last_epoch (int, optional): the index of last epoch. Default: -1.
      """
  
      def __init__(self, optimizer, end_lr, num_iter, last_epoch=-1):
          self.end_lr = end_lr
          self.num_iter = num_iter
          super(ExponentialLR, self).__init__(optimizer, last_epoch)
  
      def get_lr(self):
          curr_iter = self.last_epoch + 1
          r = curr_iter / self.num_iter
          return [base_lr * (self.end_lr / base_lr) ** r for base_lr in self.base_lrs]
  
  
  class StateCacher(object):
      def __init__(self, in_memory, cache_dir=None):
          self.in_memory = in_memory
          self.cache_dir = cache_dir
  
          if self.cache_dir is None:
              import tempfile
  
              self.cache_dir = tempfile.gettempdir()
          else:
              if not os.path.isdir(self.cache_dir):
                  raise ValueError("Given `cache_dir` is not a valid directory.")
  
          self.cached = {}
  
      def store(self, key, state_dict):
          if self.in_memory:
              self.cached.update({key: copy.deepcopy(state_dict)})
          else:
              fn = os.path.join(self.cache_dir, "state_{}_{}.pt".format(key, id(self)))
              self.cached.update({key: fn})
              torch.save(state_dict, fn)
  
      def retrieve(self, key):
          if key not in self.cached:
              raise KeyError("Target {} was not cached.".format(key))
  
          if self.in_memory:
              return self.cached.get(key)
          else:
              fn = self.cached.get(key)
              if not os.path.exists(fn):
                  raise RuntimeError(
                      "Failed to load state in {}. File doesn't exist anymore.".format(fn)
                  )
              state_dict = torch.load(fn, map_location=lambda storage, location: storage)
              return state_dict
  
      def __del__(self):
          """Check whether there are unused cached files existing in `cache_dir` before
          this instance being destroyed."""
  
          if self.in_memory:
              return
  
          for k in self.cached:
              if os.path.exists(self.cached[k]):
                  os.remove(self.cached[k])
  
  
  class DataLoaderIterWrapper(object):
      """A wrapper for iterating `torch.utils.data.DataLoader` with the ability to reset
      itself while `StopIteration` is raised."""
  
      def __init__(self, data_loader, auto_reset=True):
          self.data_loader = data_loader
          self.auto_reset = auto_reset
          self._iterator = iter(data_loader)
  
      def __next__(self):
          # Get a new set of inputs and labels
          try:
              inputs, labels = next(self._iterator)
          except StopIteration:
              if not self.auto_reset:
                  raise
              self._iterator = iter(self.data_loader)
              inputs, labels, *_ = next(self._iterator)
  
          return inputs, labels
  
      # make it compatible with python 2
      next = __next__
  
      def get_batch(self):
          return next(self)
  