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: Constructor.
  
  Args:
      model: wrapped model.
      optimizer: wrapped optimizer.
      criterion: wrapped loss function.
      device: device on which to test. run 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: 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: 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.
      amp: use Automatic Mixed Precision
      pickle_module: module used for pickling metadata and objects, default to `pickle`.
          this arg is used by `torch.save`, for more details, please check:
          https://pytorch.org/docs/stable/generated/torch.save.html#torch.save.
      pickle_protocol: can be specified to override the default protocol, default to `2`.
          this arg is used by `torch.save`, for more details, please check:
          https://pytorch.org/docs/stable/generated/torch.save.html#torch.save.
      verbose: verbose output
  Returns:
      None
code: def __init__(
      self,
      model: nn.Module,
      optimizer: Optimizer,
      criterion: torch.nn.Module,
      device: str | torch.device | None = None,
      memory_cache: bool = True,
      cache_dir: str | None = None,
      amp: bool = False,
      pickle_module: types.ModuleType = pickle,
      pickle_protocol: int = DEFAULT_PROTOCOL,
      verbose: bool = True,
  ) -> None:
      """Constructor.
  
      Args:
          model: wrapped model.
          optimizer: wrapped optimizer.
          criterion: wrapped loss function.
          device: device on which to test. run 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: 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: 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.
          amp: use Automatic Mixed Precision
          pickle_module: module used for pickling metadata and objects, default to `pickle`.
              this arg is used by `torch.save`, for more details, please check:
              https://pytorch.org/docs/stable/generated/torch.save.html#torch.save.
          pickle_protocol: can be specified to override the default protocol, default to `2`.
              this arg is used by `torch.save`, for more details, please check:
              https://pytorch.org/docs/stable/generated/torch.save.html#torch.save.
          verbose: verbose output
      Returns:
          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: dict[str, list] = {"lr": [], "loss": []}
      self.memory_cache = memory_cache
      self.cache_dir = cache_dir
      self.amp = amp
      self.verbose = verbose
  
      # 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(
          in_memory=memory_cache,
          cache_dir=cache_dir,
          pickle_module=pickle_module,
          pickle_protocol=pickle_protocol,
      )
      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
      self.device = device if device else self.model_device
  
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 1 to create a new LearningRateFinder object  with example: obj = LearningRateFinder(model, optimizer, criterion, device, memory_cache, cache_dir, amp, pickle_module, pickle_protocol, verbose)