Learning Rate finder (#1347)

* initial structure

* rebase

* incorporate suggestions

* update CHANGELOG.md

* initial docs

* fixes based on reviews

* added trainer arg

* update docs

* added saving/restore of model state

* initial tests

* fix styling

* added more tests

* fix docs, backward compatility and progressbar

* fix styling

* docs update

* updates based on review

* changed saving to standard functions

* consistent naming

* fix formatting

* improve docs, added support for nested fields, improve codecov

* update CHANGELOG.md

* Update lr_finder.rst

* Update pytorch_lightning/trainer/trainer.py

* Update trainer.py

* Update CHANGELOG.md

* Update path

* restoring

* test

* attribs

* docs

* doc typo

Co-authored-by: Nicki Skafte <nugginea@gmail.com>
Co-authored-by: William Falcon <waf2107@columbia.edu>
Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
Co-authored-by: J. Borovec <jirka.borovec@seznam.cz>

CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Added
 
 - Added `auto_select_gpus` flag to trainer that enables automatic selection of available GPUs on exclusive mode systems.
+- Added learining rate finder ([#1347](https://github.com/PyTorchLightning/pytorch-lightning/pull/1347))
 
 - 
 

docs/source/index.rst

@@ -66,6 +66,7 @@ PyTorch Lightning Documentation
    fast_training
    hooks
    hyperparameters
+   lr_finder
    multi_gpu
    multiple_loaders
    weights_loading

docs/source/lr_finder.rst

@@ -0,0 +1,108 @@
+Learning Rate Finder
+--------------------
+
+For training deep neural networks, selecting a good learning rate is essential
+for both better performance and faster convergence. Even optimizers such as
+`Adam` that are self-adjusting the learning rate can benefit from more optimal
+choices.
+
+To reduce the amount of guesswork concerning choosing a good initial learning
+rate, a `learning rate finder` can be used. As described in this `paper <https://arxiv.org/abs/1506.01186>`_ 
+a learning rate finder does a small run where the learning rate is increased 
+after each processed batch and the corresponding loss is logged. The result of 
+this is a `lr` vs. `loss` plot that can be used as guidence for choosing a optimal
+initial lr. 
+
+.. warning:: For the moment, this feature only works with models having a single optimizer.
+
+Using Lightnings build-in LR finder
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+In the most basic use case, this feature can be enabled during trainer construction
+with ``Trainer(auto_lr_find=True)``. When ``.fit(model)`` is called, the lr finder
+will automatically be run before any training is done. The ``lr`` that is found
+and used will be written to the console and logged together with all other
+hyperparameters of the model.
+    
+.. code-block:: python
+        
+    # default, no automatic learning rate finder
+    Trainer(auto_lr_find=True)
+
+When the ``lr`` or ``learning_rate`` key in hparams exists, this flag sets your learning_rate.
+In both cases, if the respective fields are not found, an error will be thrown.
+        
+.. code-block:: python
+
+    class LitModel(LightningModule):
+        def __init__(self, hparams):
+            self.hparams = hparams
+
+        def configure_optimizers(self):
+            return Adam(self.parameters(), lr=self.hparams.lr|self.hparams.learning_rate)
+
+    # finds learning rate automatically
+    # sets hparams.lr or hparams.learning_rate to that learning rate
+    Trainer(auto_lr_find=True)
+
+To use an arbitrary value set it in the parameter.
+
+.. code-block:: python
+
+    # to set to your own hparams.my_value
+    Trainer(auto_lr_find='my_value')
+
+Under the hood, when you call fit, this is what happens.  
+
+1. Run learning rate finder.
+2. Run actual fit.   
+
+.. code-block:: python
+        
+    # when you call .fit() this happens
+    # 1. find learning rate
+    # 2. actually run fit
+    trainer.fit(model)
+
+If you want to inspect the results of the learning rate finder before doing any
+actual training or just play around with the parameters of the algorithm, this
+can be done by invoking the ``lr_find`` method of the trainer. A typical example
+of this would look like
+
+.. code-block:: python
+
+    model = MyModelClass(hparams)
+    trainer = pl.Trainer()
+    
+    # Run learning rate finder
+    lr_finder = trainer.lr_find(model)
+    
+    # Results can be found in
+    lr_finder.results
+    
+    # Plot with
+    fig = lr_finder.plot(suggest=True)
+    fig.show()
+    
+    # Pick point based on plot, or get suggestion
+    new_lr = lr_finder.suggestion()
+    
+    # update hparams of the model
+    model.hparams.lr = new_lr
+    
+    # Fit model
+    trainer.fit(model)
+    
+The figure produced by ``lr_finder.plot()`` should look something like the figure
+below. It is recommended to not pick the learning rate that achives the lowest
+loss, but instead something in the middle of the sharpest downward slope (red point).
+This is the point returned py ``lr_finder.suggestion()``.
+
+.. figure:: /_images/trainer/lr_finder.png
+
+The parameters of the algorithm can be seen below.
+
+.. autoclass:: pytorch_lightning.trainer.lr_finder.TrainerLRFinderMixin
+   :members: lr_find
+   :noindex:
+   :exclude-members: _run_lr_finder_internally, save_checkpoint, restore

pytorch_lightning/trainer/__init__.py

@@ -135,6 +135,27 @@ Example::
     # default used by the Trainer
     trainer = Trainer(amp_level='O1')
 
+auto_lr_find
+^^^^^^^^^^^^
+Runs a learning rate finder algorithm (see this `paper <https://arxiv.org/abs/1506.01186>`_)
+before any training, to find optimal initial learning rate.
+
+.. code-block:: python
+
+    # default used by the Trainer (no learning rate finder)
+    trainer = Trainer(auto_lr_find=False)
+
+Example::
+
+    # run learning rate finder, results override hparams.learning_rate
+    trainer = Trainer(auto_lr_find=True)
+
+    # run learning rate finder, results override hparams.my_lr_arg
+    trainer = Trainer(auto_lr_find='my_lr_arg')
+
+.. note::
+    See the `learning rate finder guide <lr_finder.rst>`_
+
 benchmark
 ^^^^^^^^^
 

pytorch_lightning/trainer/lr_finder.py

@@ -0,0 +1,445 @@
+"""
+Trainer Learning Rate Finder
+"""
+from abc import ABC, abstractmethod
+from typing import Optional
+
+import numpy as np
+import torch
+from torch.optim.lr_scheduler import _LRScheduler
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+import os
+
+from pytorch_lightning.core.lightning import LightningModule
+from pytorch_lightning.callbacks import Callback
+from pytorch_lightning import _logger as log
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
+
+
+class TrainerLRFinderMixin(ABC):
+    @abstractmethod
+    def save_checkpoint(self, *args):
+        """Warning: this is just empty shell for code implemented in other class."""
+
+    @abstractmethod
+    def restore(self, *args):
+        """Warning: this is just empty shell for code implemented in other class."""
+
+    def _run_lr_finder_internally(self, model: LightningModule):
+        """ Call lr finder internally during Trainer.fit() """
+        lr_finder = self.lr_find(model)
+        lr = lr_finder.suggestion()
+        # TODO: log lr.results to self.logger
+        if isinstance(self.auto_lr_find, str):
+            # Try to find requested field, may be nested
+            if _nested_hasattr(model.hparams, self.auto_lr_find):
+                _nested_setattr(model.hparams, self.auto_lr_find, lr)
+            else:
+                raise MisconfigurationException(
+                    f'`auto_lr_find` was set to {self.auto_lr_find}, however'
+                    ' could not find this as a field in `model.hparams`.')
+        else:
+            if hasattr(model.hparams, 'lr'):
+                model.hparams.lr = lr
+            elif hasattr(model.hparams, 'learning_rate'):
+                model.hparams.learning_rate = lr
+            else:
+                raise MisconfigurationException(
+                    'When auto_lr_find is set to True, expects that hparams'
+                    ' either has field `lr` or `learning_rate` that can overridden')
+        log.info(f'Learning rate set to {lr}')
+
+    def lr_find(self,
+                model: LightningModule,
+                train_dataloader: Optional[DataLoader] = None,
+                min_lr: float = 1e-8,
+                max_lr: float = 1,
+                num_training: int = 100,
+                mode: str = 'exponential',
+                num_accumulation_steps: int = 1):
+        r"""
+        lr_find enables the user to do a range test of good initial learning rates,
+        to reduce the amount of guesswork in picking a good starting learning rate.
+
+        Args:
+            model: Model to do range testing for
+
+            train_dataloader: A PyTorch
+                DataLoader with training samples. If the model has
+                a predefined train_dataloader method this will be skipped.
+
+            min_lr: minimum learning rate to investigate
+
+            max_lr: maximum learning rate to investigate
+
+            num_training: number of learning rates to test
+
+            mode: search strategy, either 'linear' or 'exponential'. If set to
+                'linear' the learning rate will be searched by linearly increasing
+                after each batch. If set to 'exponential', will increase learning
+                rate exponentially.
+
+            num_accumulation_steps: number of batches to calculate loss over.
+
+        Example::
+
+            # Setup model and trainer
+            model = MyModelClass(hparams)
+            trainer = pl.Trainer()
+
+            # Run lr finder
+            lr_finder = trainer.lr_find(model, ...)
+
+            # Inspect results
+            fig = lr_finder.plot(); fig.show()
+            suggested_lr = lr_finder.suggest()
+
+            # Overwrite lr and create new model
+            hparams.lr = suggested_lr
+            model = MyModelClass(hparams)
+
+            # Ready to train with new learning rate
+            trainer.fit(model)
+
+        """
+        save_path = os.path.join(self.default_root_dir, 'lr_find_temp.ckpt')
+
+        self._dump_params(model)
+
+        # Prevent going into infinite loop
+        self.auto_lr_find = False
+
+        # Initialize lr finder object (stores results)
+        lr_finder = _LRFinder(mode, min_lr, max_lr, num_training)
+
+        # Use special lr logger callback
+        self.callbacks = [_LRCallback(num_training, show_progress_bar=True)]
+
+        # No logging
+        self.logger = None
+
+        # Max step set to number of iterations
+        self.max_steps = num_training
+
+        # Disable standard progress bar for fit
+        self.progress_bar_refresh_rate = False
+
+        # Accumulation of gradients
+        self.accumulate_grad_batches = num_accumulation_steps
+
+        # Disable standard checkpoint
+        self.checkpoint_callback = False
+
+        # Required for saving the model
+        self.optimizers, self.schedulers = [], [],
+        self.model = model
+
+        # Dump model checkpoint
+        self.save_checkpoint(str(save_path))
+
+        # Configure optimizer and scheduler
+        optimizers, _, _ = self.init_optimizers(model)
+
+        if len(optimizers) != 1:
+            raise MisconfigurationException(
+                f'`model.configure_optimizers()` returned {len(optimizers)}, but'
+                ' learning rate finder only works with single optimizer')
+        configure_optimizers = model.configure_optimizers
+        model.configure_optimizers = lr_finder._get_new_optimizer(optimizers[0])
+
+        # Fit, lr & loss logged in callback
+        self.fit(model, train_dataloader=train_dataloader)
+
+        # Prompt if we stopped early
+        if self.global_step != num_training:
+            log.info('LR finder stopped early due to diverging loss.')
+
+        # Transfer results from callback to lr finder object
+        lr_finder.results.update({'lr': self.callbacks[0].lrs,
+                                  'loss': self.callbacks[0].losses})
+
+        # Reset model state
+        self.restore(str(save_path), on_gpu=self.on_gpu)
+        os.remove(save_path)
+
+        # Finish by resetting variables so trainer is ready to fit model
+        self._restore_params(model)
+
+        return lr_finder
+
+    def _dump_params(self, model):
+        # Prevent going into infinite loop
+        self._params = {
+            'auto_lr_find': self.auto_lr_find,
+            'callbacks': self.callbacks,
+            'logger': self.logger,
+            'max_steps': self.max_steps,
+            'progress_bar_refresh_rate': self.progress_bar_refresh_rate,
+            'accumulate_grad_batches': self.accumulate_grad_batches,
+            'checkpoint_callback': self.checkpoint_callback,
+            'configure_optimizers': model.configure_optimizers,
+        }
+
+    def _restore_params(self, model):
+        self.auto_lr_find = self._params['auto_lr_find']
+        self.logger = self._params['logger']
+        self.callbacks = self._params['callbacks']
+        self.max_steps = self._params['max_steps']
+        self.progress_bar_refresh_rate = self._params['progress_bar_refresh_rate']
+        self.accumulate_grad_batches = self._params['accumulate_grad_batches']
+        self.checkpoint_callback = self._params['checkpoint_callback']
+        model.configure_optimizers = self._params['configure_optimizers']
+
+
+class _LRFinder(object):
+    """ LR finder object. This object stores the results of Trainer.lr_find().
+
+    Args:
+        mode: either `linear` or `exponential`, how to increase lr after each step
+
+        lr_min: lr to start search from
+
+        lr_max: lr to stop seach
+
+        num_training: number of steps to take between lr_min and lr_max
+
+    Example::
+        # Run lr finder
+        lr_finder = trainer.lr_find(model)
+
+        # Results stored in
+        lr_finder.results
+
+        # Plot using
+        lr_finder.plot()
+
+        # Get suggestion
+        lr = lr_finder.suggestion()
+    """
+    def __init__(self, mode: str, lr_min: float, lr_max: float, num_training: int):
+        assert mode in ('linear', 'exponential'), \
+            'mode should be either `linear` or `exponential`'
+
+        self.mode = mode
+        self.lr_min = lr_min
+        self.lr_max = lr_max
+        self.num_training = num_training
+
+        self.results = {}
+
+    def _get_new_optimizer(self, optimizer: torch.optim.Optimizer):
+        """ Construct a new `configure_optimizers()` method, that has a optimizer
+            with initial lr set to lr_min and a scheduler that will either
+            linearly or exponentially increase the lr to lr_max in num_training steps.
+
+        Args:
+            optimizer: instance of `torch.optim.Optimizer`
+
+        """
+        new_lrs = [self.lr_min] * len(optimizer.param_groups)
+        for param_group, new_lr in zip(optimizer.param_groups, new_lrs):
+            param_group["lr"] = new_lr
+            param_group["initial_lr"] = new_lr
+
+        args = (optimizer, self.lr_max, self.num_training)
+        scheduler = _LinearLR(*args) if self.mode == 'linear' else _ExponentialLR(*args)
+
+        def configure_optimizers():
+            return [optimizer], [{'scheduler': scheduler,
+                                  'interval': 'step'}]
+
+        return configure_optimizers
+
+    def plot(self, suggest: bool = False, show: bool = False):
+        """ Plot results from lr_find run
+        Args:
+            suggest: if True, will mark suggested lr to use with a red point
+
+            show: if True, will show figure
+        """
+        import matplotlib.pyplot as plt
+
+        lrs = self.results["lr"]
+        losses = self.results["loss"]
+
+        fig, ax = plt.subplots()
+
+        # Plot loss as a function of the learning rate
+        ax.plot(lrs, losses)
+        if self.mode == 'exponential':
+            ax.set_xscale("log")
+        ax.set_xlabel("Learning rate")
+        ax.set_ylabel("Loss")
+
+        if suggest:
+            _ = self.suggestion()
+            if self._optimal_idx:
+                ax.plot(lrs[self._optimal_idx], losses[self._optimal_idx],
+                        markersize=10, marker='o', color='red')
+
+        if show:
+            plt.show()
+
+        return fig
+
+    def suggestion(self):
+        """ This will propose a suggestion for choice of initial learning rate
+        as the point with the steepest negative gradient.
+
+        Returns:
+            lr: suggested initial learning rate to use
+
+        """
+        try:
+            min_grad = (np.gradient(np.array(self.results["loss"]))).argmin()
+            self._optimal_idx = min_grad
+            return self.results["lr"][min_grad]
+        except Exception:
+            log.warning('Failed to compute suggesting for `lr`.'
+                        ' There might not be enough points.')
+            self._optimal_idx = None
+
+
+class _LRCallback(Callback):
+    """ Special callback used by the learning rate finder. This callbacks log
+    the learning rate before each batch and log the corresponding loss after
+    each batch. """
+    def __init__(self, num_training: int, show_progress_bar: bool = False, beta: float = 0.98):
+        self.num_training = num_training
+        self.beta = beta
+        self.losses = []
+        self.lrs = []
+        self.avg_loss = 0.0
+        self.best_loss = 0.0
+        self.show_progress_bar = show_progress_bar
+        self.progress_bar = None
+
+    def on_batch_start(self, trainer, pl_module):
+        """ Called before each training batch, logs the lr that will be used """
+        if self.show_progress_bar and self.progress_bar is None:
+            self.progress_bar = tqdm(desc='Finding best initial lr', total=self.num_training)
+
+        self.lrs.append(trainer.lr_schedulers[0]['scheduler'].lr[0])
+
+    def on_batch_end(self, trainer, pl_module):
+        """ Called when the training batch ends, logs the calculated loss """
+        if self.progress_bar:
+            self.progress_bar.update()
+
+        current_loss = trainer.running_loss.last().item()
+        current_step = trainer.global_step + 1  # remove the +1 in 1.0
+
+        # Avg loss (loss with momentum) + smoothing
+        self.avg_loss = self.beta * self.avg_loss + (1 - self.beta) * current_loss
+        smoothed_loss = self.avg_loss / (1 - self.beta**current_step)
+
+        # Check if we diverging
+        if current_step > 1 and smoothed_loss > 4 * self.best_loss:
+            trainer.max_steps = current_step  # stop signal
+            if self.progress_bar:
+                self.progress_bar.close()
+
+        # Save best loss for diverging checking
+        if smoothed_loss < self.best_loss or current_step == 1:
+            self.best_loss = smoothed_loss
+
+        self.losses.append(smoothed_loss)
+
+
+class _LinearLR(_LRScheduler):
+    """Linearly increases the learning rate between two boundaries
+    over a number of iterations.
+    Arguments:
+
+        optimizer: wrapped optimizer.
+
+        end_lr: the final learning rate.
+
+        num_iter: the number of iterations over which the test occurs.
+
+        last_epoch: the index of last epoch. Default: -1.
+    """
+
+    def __init__(self,
+                 optimizer: torch.optim.Optimizer,
+                 end_lr: float,
+                 num_iter: int,
+                 last_epoch: int = -1):
+        self.end_lr = end_lr
+        self.num_iter = num_iter
+        super(_LinearLR, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        curr_iter = self.last_epoch + 1
+        r = curr_iter / self.num_iter
+
+        if self.last_epoch > 0:
+            val = [base_lr + r * (self.end_lr - base_lr) for base_lr in self.base_lrs]
+        else:
+            val = [base_lr for base_lr in self.base_lrs]
+        self._lr = val
+        return val
+
+    @property
+    def lr(self):
+        return self._lr
+
+
+class _ExponentialLR(_LRScheduler):
+    """Exponentially increases the learning rate between two boundaries
+    over a number of iterations.
+
+    Arguments:
+
+        optimizer: wrapped optimizer.
+
+        end_lr: the final learning rate.
+
+        num_iter: the number of iterations over which the test occurs.
+
+        last_epoch: the index of last epoch. Default: -1.
+    """
+
+    def __init__(self,
+                 optimizer: torch.optim.Optimizer,
+                 end_lr: float,
+                 num_iter: int,
+                 last_epoch: int = -1):
+        self.end_lr = end_lr
+        self.num_iter = num_iter
+        super(_ExponentialLR, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        curr_iter = self.last_epoch + 1
+        r = curr_iter / self.num_iter
+
+        if self.last_epoch > 0:
+            val = [base_lr * (self.end_lr / base_lr) ** r for base_lr in self.base_lrs]
+        else:
+            val = [base_lr for base_lr in self.base_lrs]
+        self._lr = val
+        return val
+
+    @property
+    def lr(self):
+        return self._lr
+
+
+def _nested_hasattr(obj, path):
+    parts = path.split(".")
+    for part in parts:
+        if hasattr(obj, part):
+            obj = getattr(obj, part)
+        else:
+            return False
+    else:
+        return True
+
+
+def _nested_setattr(obj, path, val):
+    parts = path.split(".")
+    for part in parts[:-1]:
+        if hasattr(obj, part):
+            obj = getattr(obj, part)
+    setattr(obj, parts[-1], val)

pytorch_lightning/trainer/trainer.py

@@ -36,9 +36,11 @@ from pytorch_lightning.trainer.supporters import TensorRunningAccum
 from pytorch_lightning.trainer.training_io import TrainerIOMixin
 from pytorch_lightning.trainer.training_loop import TrainerTrainLoopMixin
 from pytorch_lightning.trainer.training_tricks import TrainerTrainingTricksMixin
+from pytorch_lightning.trainer.lr_finder import TrainerLRFinderMixin
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities import rank_zero_warn
 
+
 try:
     from apex import amp
 except ImportError:
@@ -70,6 +72,7 @@ class Trainer(
     TrainerTrainLoopMixin,
     TrainerCallbackConfigMixin,
     TrainerCallbackHookMixin,
+    TrainerLRFinderMixin,
     TrainerDeprecatedAPITillVer0_8,
     TrainerDeprecatedAPITillVer0_9,
 ):
@@ -122,6 +125,7 @@ class Trainer(
             profiler: Optional[BaseProfiler] = None,
             benchmark: bool = False,
             reload_dataloaders_every_epoch: bool = False,
+            auto_lr_find: Union[bool, str] = False,
             default_save_path=None,  # backward compatible, todo: remove in v0.8.0
             gradient_clip=None,  # backward compatible, todo: remove in v0.8.0
             nb_gpu_nodes=None,  # backward compatible, todo: remove in v0.8.0
@@ -271,6 +275,11 @@ class Trainer(
 
             reload_dataloaders_every_epoch: Set to True to reload dataloaders every epoch
 
+            auto_lr_find: If set to True, will `initially` run a learning rate finder,
+                trying to optimize initial learning for faster convergence. Sets learning
+                rate in self.hparams.lr | self.hparams.learning_rate in the lightning module.
+                To use a different key, set a string instead of True with the key name.
+
             benchmark: If true enables cudnn.benchmark.
         """
 
@@ -344,6 +353,8 @@ class Trainer(
 
         self.reload_dataloaders_every_epoch = reload_dataloaders_every_epoch
 
+        self.auto_lr_find = auto_lr_find
+
         self.truncated_bptt_steps = truncated_bptt_steps
         self.resume_from_checkpoint = resume_from_checkpoint
         self.shown_warnings = set()
@@ -696,6 +707,10 @@ class Trainer(
         # only on proc 0 because no spawn has happened yet
         model.prepare_data()
 
+        # Run learning rate finder:
+        if self.auto_lr_find:
+            self._run_lr_finder_internally(model)
+
         # route to appropriate start method
         # when using multi-node or DDP within a node start each module in a separate process
         if self.use_ddp2:

requirements-extra.txt

@@ -6,3 +6,4 @@ mlflow>=1.0.0
 test_tube>=0.7.5
 wandb>=0.8.21
 trains>=0.14.1
+matplotlib>=3.1.1

tests/trainer/test_lr_finder.py

@@ -0,0 +1,181 @@
+import pytest
+
+import torch
+import tests.base.utils as tutils
+from pytorch_lightning import Trainer
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from tests.base import (
+    LightTrainDataloader,
+    TestModelBase,
+    LightTestMultipleOptimizersWithSchedulingMixin,
+)
+
+
+def test_error_on_more_than_1_optimizer(tmpdir):
+    ''' Check that error is thrown when more than 1 optimizer is passed '''
+    tutils.reset_seed()
+
+    class CurrentTestModel(
+        LightTestMultipleOptimizersWithSchedulingMixin,
+        LightTrainDataloader,
+        TestModelBase,
+    ):
+        pass
+
+    hparams = tutils.get_default_hparams()
+    model = CurrentTestModel(hparams)
+
+    # logger file to get meta
+    trainer = Trainer(
+        default_save_path=tmpdir,
+        max_epochs=1
+    )
+
+    with pytest.raises(MisconfigurationException):
+        trainer.lr_find(model)
+
+
+def test_model_reset_correctly(tmpdir):
+    ''' Check that model weights are correctly reset after lr_find() '''
+    tutils.reset_seed()
+
+    class CurrentTestModel(
+        LightTrainDataloader,
+        TestModelBase,
+    ):
+        pass
+
+    hparams = tutils.get_default_hparams()
+    model = CurrentTestModel(hparams)
+
+    # logger file to get meta
+    trainer = Trainer(
+        default_save_path=tmpdir,
+        max_epochs=1
+    )
+
+    before_state_dict = model.state_dict()
+
+    _ = trainer.lr_find(model, num_training=5)
+
+    after_state_dict = model.state_dict()
+
+    for key in before_state_dict.keys():
+        assert torch.all(torch.eq(before_state_dict[key], after_state_dict[key])), \
+            'Model was not reset correctly after learning rate finder'
+
+
+def test_trainer_reset_correctly(tmpdir):
+    ''' Check that all trainer parameters are reset correctly after lr_find() '''
+    tutils.reset_seed()
+
+    class CurrentTestModel(
+        LightTrainDataloader,
+        TestModelBase,
+    ):
+        pass
+
+    hparams = tutils.get_default_hparams()
+    model = CurrentTestModel(hparams)
+
+    # logger file to get meta
+    trainer = Trainer(
+        default_save_path=tmpdir,
+        max_epochs=1
+    )
+
+    changed_attributes = ['callbacks', 'logger', 'max_steps', 'auto_lr_find',
+                          'progress_bar_refresh_rate',
+                          'accumulate_grad_batches',
+                          'checkpoint_callback']
+    attributes_before = {}
+    for ca in changed_attributes:
+        attributes_before[ca] = getattr(trainer, ca)
+
+    _ = trainer.lr_find(model, num_training=5)
+
+    attributes_after = {}
+    for ca in changed_attributes:
+        attributes_after[ca] = getattr(trainer, ca)
+
+    for key in changed_attributes:
+        assert attributes_before[key] == attributes_after[key], \
+            f'Attribute {key} was not reset correctly after learning rate finder'
+
+
+def test_trainer_arg_bool(tmpdir):
+    tutils.reset_seed()
+
+    class CurrentTestModel(
+        LightTrainDataloader,
+        TestModelBase,
+    ):
+        pass
+
+    hparams = tutils.get_default_hparams()
+    model = CurrentTestModel(hparams)
+    before_lr = hparams.learning_rate
+    # logger file to get meta
+    trainer = Trainer(
+        default_save_path=tmpdir,
+        max_epochs=1,
+        auto_lr_find=True
+    )
+
+    trainer.fit(model)
+    after_lr = model.hparams.learning_rate
+    assert before_lr != after_lr, \
+        'Learning rate was not altered after running learning rate finder'
+
+
+def test_trainer_arg_str(tmpdir):
+    tutils.reset_seed()
+
+    class CurrentTestModel(
+        LightTrainDataloader,
+        TestModelBase,
+    ):
+        pass
+
+    hparams = tutils.get_default_hparams()
+    hparams.__dict__['my_fancy_lr'] = 1.0  # update with non-standard field
+    model = CurrentTestModel(hparams)
+    before_lr = hparams.my_fancy_lr
+    # logger file to get meta
+    trainer = Trainer(
+        default_save_path=tmpdir,
+        max_epochs=1,
+        auto_lr_find='my_fancy_lr'
+    )
+
+    trainer.fit(model)
+    after_lr = model.hparams.my_fancy_lr
+    assert before_lr != after_lr, \
+        'Learning rate was not altered after running learning rate finder'
+
+
+def test_call_to_trainer_method(tmpdir):
+    tutils.reset_seed()
+
+    class CurrentTestModel(
+        LightTrainDataloader,
+        TestModelBase,
+    ):
+        pass
+
+    hparams = tutils.get_default_hparams()
+    model = CurrentTestModel(hparams)
+    before_lr = hparams.learning_rate
+    # logger file to get meta
+    trainer = Trainer(
+        default_save_path=tmpdir,
+        max_epochs=1,
+    )
+
+    lrfinder = trainer.lr_find(model, mode='linear')
+    after_lr = lrfinder.suggestion()
+    model.hparams.learning_rate = after_lr
+    trainer.fit(model)
+
+    assert before_lr != after_lr, \
+        'Learning rate was not altered after running learning rate finder'