Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/tuner/lr_finder.py (Python)

Code before:
from torch.optim.lr_scheduler import _LRScheduler
from torch.utils.data import DataLoader

from pytorch_lightning import _logger as log
from pytorch_lightning.callbacks import Callback
from pytorch_lightning.core.datamodule import LightningDataModule
from pytorch_lightning.core.lightning import LightningModule
from pytorch_lightning.loggers.base import DummyLogger
from pytorch_lightning.utilities.exceptions import MisconfigurationException
from pytorch_lightning.utilities.parsing import lightning_hasattr, lightning_setattr
from pytorch_lightning.utilities import rank_zero_warn

# check if ipywidgets is installed before importing tqdm.auto
# to ensure it won't fail and a progress bar is displayed
if importlib.util.find_spec('ipywidgets') is not None:
    from tqdm.auto import tqdm
else:
    from tqdm import tqdm


def _run_lr_finder_internally(trainer, model: LightningModule):
    """ Call lr finder internally during Trainer.fit() """
    if trainer.fast_dev_run:
        rank_zero_warn('Skipping learning rate finder since `fast_dev_run=True`', UserWarning)
    lr_finder = lr_find(trainer, model)
    lr = lr_finder.suggestion()

    # TODO: log lr.results to self.logger
    if isinstance(trainer.auto_lr_find, str):
        # Try to find requested field, may be nested
        if lightning_hasattr(model, trainer.auto_lr_find):
            lightning_setattr(model, trainer.auto_lr_find, lr)
        else:
            raise MisconfigurationException(
                f'`auto_lr_find` was set to {trainer.auto_lr_find}, however'
                ' could not find this as a field in `model` or `model.hparams`.')
    else:
        if lightning_hasattr(model, 'lr'):
            lightning_setattr(model, 'lr', lr)
        elif lightning_hasattr(model, 'learning_rate'):
            lightning_setattr(model, 'learning_rate', lr)
        else:
            raise MisconfigurationException(
                'When auto_lr_find is set to True, expects that `model` or'
                ' `model.hparams` either has field `lr` or `learning_rate`'
                ' that can overridden')
    log.info(f'Learning rate set to {lr}')


def lr_find(
        trainer,

Review comment (performance):
this will still call the `lr_find` right? Why not keep this condition under inside the functions `lr_find` since they are both calling it and the same for `scale_batch_size`?

Code after:
from torch.optim.lr_scheduler import _LRScheduler
from torch.utils.data import DataLoader

from pytorch_lightning import _logger as log
from pytorch_lightning.callbacks import Callback
from pytorch_lightning.core.datamodule import LightningDataModule
from pytorch_lightning.core.lightning import LightningModule
from pytorch_lightning.loggers.base import DummyLogger
from pytorch_lightning.utilities.exceptions import MisconfigurationException
from pytorch_lightning.utilities.parsing import lightning_hasattr, lightning_setattr
from pytorch_lightning.utilities import rank_zero_warn
from pytorch_lightning.utilities.cloud_io import get_filesystem

# check if ipywidgets is installed before importing tqdm.auto
# to ensure it won't fail and a progress bar is displayed
if importlib.util.find_spec('ipywidgets') is not None:
    from tqdm.auto import tqdm
else:
    from tqdm import tqdm


def _run_lr_finder_internally(trainer, model: LightningModule):
    """ Call lr finder internally during Trainer.fit() """
    lr_finder = lr_find(trainer, model)

    if lr_finder is None:
        return

    lr = lr_finder.suggestion()

    # TODO: log lr.results to self.logger
    if isinstance(trainer.auto_lr_find, str):
        # Try to find requested field, may be nested
        if lightning_hasattr(model, trainer.auto_lr_find):
            lightning_setattr(model, trainer.auto_lr_find, lr)
        else:
            raise MisconfigurationException(
                f'`auto_lr_find` was set to {trainer.auto_lr_find}, however'
                ' could not find this as a field in `model` or `model.hparams`.')
    else:
        if lightning_hasattr(model, 'lr'):
            lightning_setattr(model, 'lr', lr)
        elif lightning_hasattr(model, 'learning_rate'):
            lightning_setattr(model, 'learning_rate', lr)
        else:
            raise MisconfigurationException(
                'When auto_lr_find is set to True, expects that `model` or'
                ' `model.hparams` either has field `lr` or `learning_rate`'
                ' that can overridden')
    log.info(f'Learning rate set to {lr}')