class_docstr: The ``BatchSizeFinder`` callback tries to find the largest batch size for a given model that does not give an
  out of memory (OOM) error. All you need to do is add it as a callback inside Trainer and call
  ``trainer.{fit,validate,test,predict}``. Internally it calls the respective step function ``steps_per_trial`` times
  for each batch size until one of the batch sizes generates an OOM error.
  
  .. warning::  This is an :ref:`experimental <versioning:Experimental API>` feature.
  
  Args:
      mode: search strategy to update the batch size:
  
          - ``'power'``: Keep multiplying the batch size by 2, until we get an OOM error.
          - ``'binsearch'``: Initially keep multiplying by 2 and after encountering an OOM error
            do a binary search between the last successful batch size and the batch size that failed.
  
      steps_per_trial: number of steps to run with a given batch size.
          Ideally 1 should be enough to test if an OOM error occurs,
          however in practice a few are needed.
  
      init_val: initial batch size to start the search with.
  
      max_trials: max number of increases in batch size done before
          algorithm is terminated
  
      batch_arg_name: name of the attribute that stores the batch size.
          It is expected that the user has provided a model or datamodule that has a hyperparameter
          with that name. We will look for this attribute name in the following places
  
          - ``model``
          - ``model.hparams``
          - ``trainer.datamodule`` (the datamodule passed to the tune method)
  
  Example::
  
      # 1. Customize the BatchSizeFinder callback to run at different epochs. This feature is
      # useful while fine-tuning models since you can't always use the same batch size after
      # unfreezing the backbone.
      from lightning.pytorch.callbacks import BatchSizeFinder
  
  
      class FineTuneBatchSizeFinder(BatchSizeFinder):
          def __init__(self, milestones, *args, **kwargs):
              super().__init__(*args, **kwargs)
              self.milestones = milestones
  
          def on_fit_start(self, *args, **kwargs):
              return
  
          def on_train_epoch_start(self, trainer, pl_module):
              if trainer.current_epoch in self.milestones or trainer.current_epoch == 0:
                  self.scale_batch_size(trainer, pl_module)
  
  
      trainer = Trainer(callbacks=[FineTuneBatchSizeFinder(milestones=(5, 10))])
      trainer.fit(...)
  
  Example::
  
      # 2. Run batch size finder for validate/test/predict.
      from lightning.pytorch.callbacks import BatchSizeFinder
  
  
      class EvalBatchSizeFinder(BatchSizeFinder):
          def __init__(self, *args, **kwargs):
              super().__init__(*args, **kwargs)
  
          def on_fit_start(self, *args, **kwargs):
              return
  
          def on_test_start(self, trainer, pl_module):
              self.scale_batch_size(trainer, pl_module)
  
  
      trainer = Trainer(callbacks=[EvalBatchSizeFinder()])
      trainer.test(...)
class_docstr_tok: ["The","``","BatchSizeFinder","``","callback","tries","to","find","the","largest","batch","size","for","a","given","model","that","does","not","give","an","out","of","memory","(","OOM",")","error",".","All","you","need","to","do","is","add","it","as","a","callback","inside","Trainer","and","call","``","trainer.","{","fit",",","validate",",","test",",","predict","}","``",".","Internally","it","calls","the","respective","step","function","``","steps_per_trial","``","times","for","each","batch","size","until","one","of","the","batch","sizes","generates","an","OOM","error",".","..","warning",":",":","This","is","an",":","ref",":","`","experimental","<","versioning",":","Experimental","API",">","`","feature",".","Args",":","mode",":","search","strategy","to","update","the","batch","size",":","-","``","'power","'","``",":","Keep","multiplying","the","batch","size","by","2",",","until","we","get","an","OOM","error",".","-","``","'binsearch","'","``",":","Initially","keep","multiplying","by","2","and","after","encountering","an","OOM","error","do","a","binary","search","between","the","last","successful","batch","size","and","the","batch","size","that","failed",".","steps_per_trial",":","number","of","steps","to","run","with","a","given","batch","size",".","Ideally","1","should","be","enough","to","test","if","an","OOM","error","occurs",",","however","in","practice","a","few","are","needed",".","init_val",":","initial","batch","size","to","start","the","search","with",".","max_trials",":","max","number","of","increases","in","batch","size","done","before","algorithm","is","terminated","batch_arg_name",":","name","of","the","attribute","that","stores","the","batch","size",".","It","is","expected","that","the","user","has","provided","a","model","or","datamodule","that","has","a","hyperparameter","with","that","name",".","We","will","look","for","this","attribute","name","in","the","following","places","-","``","model","``","-","``","model.hparams","``","-","``","trainer.datamodule","``","(","the","datamodule","passed","to","the","tune","method",")","Example",":",":","#","1",".","Customize","the","BatchSizeFinder","callback","to","run","at","different","epochs",".","This","feature","is","#","useful","while","fine-tuning","models","since","you","ca","n't","always","use","the","same","batch","size","after","#","unfreezing","the","backbone",".","from","lightning.pytorch.callbacks","import","BatchSizeFinder","class","FineTuneBatchSizeFinder","(","BatchSizeFinder",")",":","def","__init__","(","self",",","milestones",",","*","args",",","*","*","kwargs",")",":","super","(",")",".__init__","(","*","args",",","*","*","kwargs",")","self.milestones","=","milestones","def","on_fit_start","(","self",",","*","args",",","*","*","kwargs",")",":","return","def","on_train_epoch_start","(","self",",","trainer",",","pl_module",")",":","if","trainer.current_epoch","in","self.milestones","or","trainer.current_epoch","==","0",":","self.scale_batch_size","(","trainer",",","pl_module",")","trainer","=","Trainer","(","callbacks=","[","FineTuneBatchSizeFinder","(","milestones=","(","5",",","10",")",")","]",")","trainer.fit","(","...",")","Example",":",":","#","2",".","Run","batch","size","finder","for","validate\/test\/predict",".","from","lightning.pytorch.callbacks","import","BatchSizeFinder","class","EvalBatchSizeFinder","(","BatchSizeFinder",")",":","def","__init__","(","self",",","*","args",",","*","*","kwargs",")",":","super","(",")",".__init__","(","*","args",",","*","*","kwargs",")","def","on_fit_start","(","self",",","*","args",",","*","*","kwargs",")",":","return","def","on_test_start","(","self",",","trainer",",","pl_module",")",":","self.scale_batch_size","(","trainer",",","pl_module",")","trainer","=","Trainer","(","callbacks=","[","EvalBatchSizeFinder","(",")","]",")","trainer.test","(","...",")"]
code: def scale_batch_size(
      self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
  ) -> None:
      new_size = _scale_batch_size(
          trainer,
          self._mode,
          self._steps_per_trial,
          self._init_val,
          self._max_trials,
          self._batch_arg_name,
      )
  
      self.optimal_batch_size = new_size
      if self._early_exit:
          raise _TunerExitException()
  
file_path: pytorch-lightning/src/lightning/pytorch/callbacks/batch_size_finder.py
imports: from typing import Optional
  from typing_extensions import override
  import lightning.pytorch
  from lightning.pytorch.callbacks.callback import Callback
  from lightning.pytorch.tuner.batch_size_scaling import _scale_batch_size
  from lightning.pytorch.utilities.exceptions import MisconfigurationException, _TunerExitException
  from lightning.pytorch.utilities.parsing import lightning_hasattr
  from lightning.pytorch.utilities.rank_zero import rank_zero_warn
label_desc: Use image node_id 3 for calling the BatchSizeFinder obj's underlying member method code with example usage: obj.scale_batch_size(trainer, pl_module) without return types