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"]
code: def _train_batch(
      self,
      train_iter: TrainDataLoaderIter,
      accumulation_steps: int,
      non_blocking_transfer: bool = True,
  ) -> float:
      self.model.train()
      total_loss = 0
  
      self.optimizer.zero_grad()
      for i in range(accumulation_steps):
          inputs, labels = next(train_iter)
          inputs, labels = copy_to_device(
              [inputs, labels],
              device=self.device,
              non_blocking=non_blocking_transfer,
          )
  
          # Forward pass
          outputs = self.model(inputs)
          loss = self.criterion(outputs, labels)
  
          # Loss should be averaged in each step
          loss /= accumulation_steps
  
          # Backward pass
          if self.amp 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 torch.cuda.amp.scale_loss(loss, self.optimizer, delay_unscale=delay_unscale) as scaled_loss:  # type: ignore
                  scaled_loss.backward()
          else:
              loss.backward()
  
          total_loss += loss.item()
  
      self.optimizer.step()
  
      return total_loss
  
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 6 for calling the LearningRateFinder obj's underlying member method code with example usage: obj._train_batch(train_iter, accumulation_steps, non_blocking_transfer) and returns: total_loss