class_docstr: 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.
  
  Example (fastai approach):
  >>> lr_finder = LearningRateFinder(net, optimizer, criterion)
  >>> lr_finder.range_test(data_loader, end_lr=100, num_iter=100)
  >>> lr_finder.get_steepest_gradient()
  >>> lr_finder.plot() # to inspect the loss-learning rate graph
  
  Example (Leslie Smith's approach):
  >>> lr_finder = LearningRateFinder(net, optimizer, criterion)
  >>> lr_finder.range_test(train_loader, 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
  >>> data_loader = torch.utils.data.DataLoader(train_data, batch_size=real_bs, shuffle=True)
  >>> acc_lr_finder = LearningRateFinder(net, optimizer, criterion)
  >>> acc_lr_finder.range_test(data_loader, end_lr=10, num_iter=100, accumulation_steps=accumulation_steps)
  
  By default, image will be extracted from data loader with x["image"] and x[0], depending on whether
  batch data is a dictionary or not (and similar behaviour for extracting the label). If your data loader
  returns something other than this, pass a callable function to extract it, e.g.:
  >>> image_extractor = lambda x: x["input"]
  >>> label_extractor = lambda x: x[100]
  >>> lr_finder = LearningRateFinder(net, optimizer, criterion)
  >>> lr_finder.range_test(train_loader, val_loader, image_extractor, label_extractor)
  
  References:
  Modified from: https://github.com/davidtvs/pytorch-lr-finder.
  Cyclical Learning Rates for Training Neural Networks: https://arxiv.org/abs/1506.01186
class_docstr_tok: ["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",".","Example","(","fastai","approach",")",":",">",">",">","lr_finder","=","LearningRateFinder","(","net",",","optimizer",",","criterion",")",">",">",">","lr_finder.range_test","(","data_loader",",","end_lr=100",",","num_iter=100",")",">",">",">","lr_finder.get_steepest_gradient","(",")",">",">",">","lr_finder.plot","(",")","#","to","inspect","the","loss-learning","rate","graph","Example","(","Leslie","Smith","'s","approach",")",":",">",">",">","lr_finder","=","LearningRateFinder","(","net",",","optimizer",",","criterion",")",">",">",">","lr_finder.range_test","(","train_loader",",","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",">",">",">","data_loader","=","torch.utils.data.DataLoader","(","train_data",",","batch_size=real_bs",",","shuffle=True",")",">",">",">","acc_lr_finder","=","LearningRateFinder","(","net",",","optimizer",",","criterion",")",">",">",">","acc_lr_finder.range_test","(","data_loader",",","end_lr=10",",","num_iter=100",",","accumulation_steps=accumulation_steps",")","By","default",",","image","will","be","extracted","from","data","loader","with","x","[","``","image","''","]","and","x","[","0","]",",","depending","on","whether","batch","data","is","a","dictionary","or","not","(","and","similar","behaviour","for","extracting","the","label",")",".","If","your","data","loader","returns","something","other","than","this",",","pass","a","callable","function","to","extract","it",",","e.g",".",":",">",">",">","image_extractor","=","lambda","x",":","x","[","``","input","''","]",">",">",">","label_extractor","=","lambda","x",":","x","[","100","]",">",">",">","lr_finder","=","LearningRateFinder","(","net",",","optimizer",",","criterion",")",">",">",">","lr_finder.range_test","(","train_loader",",","val_loader",",","image_extractor",",","label_extractor",")","References",":","Modified","from",":","https",":","\/\/github.com\/davidtvs\/pytorch-lr-finder",".","Cyclical","Learning","Rates","for","Training","Neural","Networks",":","https",":","\/\/arxiv.org\/abs\/1506.01186"]
docstr: Performs the learning rate range test.
  
  Args:
      train_loader: training set data loader.
      val_loader: validation data loader (if desired).
      image_extractor: callable function to get the image from a batch of data.
          Default: `x["image"] if isinstance(x, dict) else x[0]`.
      label_extractor: callable function to get the label from a batch of data.
          Default: `x["label"] if isinstance(x, dict) else x[1]`.
      start_lr : the starting learning rate for the range test.
          The default is the optimizer's learning rate.
      end_lr: the maximum learning rate to test. The test may stop earlier than
          this if the result starts diverging.
      num_iter: the max number of iterations for test.
      step_mode: schedule for increasing learning rate: (`linear` or `exp`).
      smooth_f: the loss smoothing factor within the `[0, 1[` interval. Disabled
          if set to `0`, otherwise loss is smoothed using exponential smoothing.
      diverge_th: test is stopped when loss surpasses threshold:
          `diverge_th * best_loss`.
      accumulation_steps: steps for gradient accumulation. If set to `1`,
          gradients are not accumulated.
      non_blocking_transfer: when `True`, moves data to device asynchronously if
          possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.
      auto_reset: if `True`, returns model and optimizer to original states at end
          of test.
  Returns:
      None
code: def range_test(
      self,
      train_loader: DataLoader,
      val_loader: DataLoader | None = None,
      image_extractor: Callable = default_image_extractor,
      label_extractor: Callable = default_label_extractor,
      start_lr: float | None = None,
      end_lr: float = 10.0,
      num_iter: int = 100,
      step_mode: str = "exp",
      smooth_f: float = 0.05,
      diverge_th: int = 5,
      accumulation_steps: int = 1,
      non_blocking_transfer: bool = True,
      auto_reset: bool = True,
  ) -> None:
      """Performs the learning rate range test.
  
      Args:
          train_loader: training set data loader.
          val_loader: validation data loader (if desired).
          image_extractor: callable function to get the image from a batch of data.
              Default: `x["image"] if isinstance(x, dict) else x[0]`.
          label_extractor: callable function to get the label from a batch of data.
              Default: `x["label"] if isinstance(x, dict) else x[1]`.
          start_lr : the starting learning rate for the range test.
              The default is the optimizer's learning rate.
          end_lr: the maximum learning rate to test. The test may stop earlier than
              this if the result starts diverging.
          num_iter: the max number of iterations for test.
          step_mode: schedule for increasing learning rate: (`linear` or `exp`).
          smooth_f: the loss smoothing factor within the `[0, 1[` interval. Disabled
              if set to `0`, otherwise loss is smoothed using exponential smoothing.
          diverge_th: test is stopped when loss surpasses threshold:
              `diverge_th * best_loss`.
          accumulation_steps: steps for gradient accumulation. If set to `1`,
              gradients are not accumulated.
          non_blocking_transfer: when `True`, moves data to device asynchronously if
              possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.
          auto_reset: if `True`, returns model and optimizer to original states at end
              of test.
      Returns:
          None
      """
  
      # Reset test results
      self.history = {"lr": [], "loss": []}
      best_loss = -float("inf")
  
      # 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)
  
      # Check number of iterations
      if num_iter <= 1:
          raise ValueError("`num_iter` must be larger than 1")
  
      # Initialize the proper learning rate policy
      lr_schedule: ExponentialLR | LinearLR
      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(
              f"expected one of (exp, linear), got {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
      train_iter = TrainDataLoaderIter(
          train_loader, image_extractor, label_extractor
      )
      if val_loader:
          val_iter = ValDataLoaderIter(
              val_loader, image_extractor, label_extractor
          )
  
      trange: partial[tqdm.trange] | type[range]
      if self.verbose and has_tqdm:
          trange = partial(
              tqdm.trange, desc="Computing optimal learning rate"
          )
          tprint = tqdm.tqdm.write
      else:
          trange = range
          tprint = print
  
      for iteration in trange(num_iter):
          if self.verbose and not has_tqdm:
              print(
                  f"Computing optimal learning rate, iteration {iteration + 1}/{num_iter}"
              )
  
          # Train on batch and retrieve loss
          loss = self._train_batch(
              train_iter,
              accumulation_steps,
              non_blocking_transfer=non_blocking_transfer,
          )
          if val_loader:
              loss = self._validate(
                  val_iter, non_blocking_transfer=non_blocking_transfer
              )
  
          # Update the learning rate
          self.history["lr"].append(lr_schedule.get_lr()[0])
          lr_schedule.step()
  
          # Track the best loss and smooth it if smooth_f is specified
          if iteration == 0:
              best_loss = loss
          else:
              if smooth_f > 0:
                  loss = (
                      smooth_f * loss
                      + (1 - smooth_f) * self.history["loss"][-1]
                  )
              if loss < best_loss:
                  best_loss = loss
  
          # Check if the loss has diverged; if it has, stop the test
          self.history["loss"].append(loss)
          if loss > diverge_th * best_loss:
              if self.verbose:
                  tprint("Stopping early, the loss has diverged")
              break
  
      if auto_reset:
          if self.verbose:
              print("Resetting model and optimizer")
          self.reset()
  
imports: from __future__ import annotations
  import pickle
  import types
  import warnings
  from functools import partial
  from typing import TYPE_CHECKING, Any, Callable
  import numpy
  import torch
  import torch.nn
  from torch.optim import Optimizer
  from torch.serialization import DEFAULT_PROTOCOL
  from torch.utils.data import DataLoader
  from monai.networks.utils import eval_mode
  from monai.optimizers.lr_scheduler import ExponentialLR, LinearLR
  from monai.utils import StateCacher, copy_to_device, optional_import
label_desc: Use image node_id 3 for calling the LearningRateFinder obj's underlying member method code with example usage: obj.range_test(train_loader, val_loader, image_extractor, label_extractor, start_lr, end_lr, num_iter, step_mode, smooth_f, diverge_th, accumulation_steps, non_blocking_transfer, auto_reset) without return types