Skip tuner algorithms on fast dev (#3903)

* skip on fast dev

* fix error

* changelog

* fix recursive issue

* combine tests

* pep8

* move logic to base funcs

* fix mistake

* Update pytorch_lightning/tuner/lr_finder.py

Co-authored-by: Rohit Gupta <[email protected]>

* pep

Co-authored-by: William Falcon <[email protected]>
Co-authored-by: Nicki Skafte <[email protected]>
Co-authored-by: Rohit Gupta <[email protected]>
Co-authored-by: chaton <[email protected]>

(cherry picked from commit 4f3160b)

Author: SkafteNicki

CHANGELOG.md

@@ -30,12 +30,12 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `fsspec` to tuner ([#4458](https://github.com/PyTorchLightning/pytorch-lightning/pull/4458))
 
 
-- Added metrics aggregation in Horovod and fixed early stopping ([#3775](https://github.com/PyTorchLightning/pytorch-lightning/pull/3775)) 
+- Added metrics aggregation in Horovod and fixed early stopping ([#3775](https://github.com/PyTorchLightning/pytorch-lightning/pull/3775))
 
 
 ### Changed
 
-
+- Tuner algorithms will be skipped if `fast_dev_run=True` ([#3903](https://github.com/PyTorchLightning/pytorch-lightning/pull/3903))
 
 ### Deprecated
 

pytorch_lightning/tuner/batch_size_scaling.py

@@ -68,6 +68,10 @@ def scale_batch_size(trainer,
         **fit_kwargs: remaining arguments to be passed to .fit(), e.g., dataloader
             or datamodule.
     """
+    if trainer.fast_dev_run:
+        rank_zero_warn('Skipping batch size scaler since `fast_dev_run=True`', UserWarning)
+        return
+
     if not lightning_hasattr(model, batch_arg_name):
         raise MisconfigurationException(
             f'Field {batch_arg_name} not found in both `model` and `model.hparams`')

pytorch_lightning/tuner/lr_finder.py

@@ -29,6 +29,8 @@
 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
@@ -41,6 +43,10 @@
 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
@@ -130,7 +136,11 @@ def lr_find(
         trainer.fit(model)
 
     """
-    save_path = os.path.join(trainer.default_root_dir, 'lr_find_temp.ckpt')
+    if trainer.fast_dev_run:
+        rank_zero_warn('Skipping learning rate finder since `fast_dev_run=True`', UserWarning)
+        return
+
+    save_path = os.path.join(trainer.default_root_dir, 'lr_find_temp_model.ckpt')
 
     __lr_finder_dump_params(trainer, model)
 
@@ -181,8 +191,11 @@ def lr_find(
     lr_finder._total_batch_idx = trainer.total_batch_idx  # for debug purpose
 
     # Reset model state
-    trainer.checkpoint_connector.restore(str(save_path), on_gpu=trainer.on_gpu)
-    os.remove(save_path)
+    if trainer.is_global_zero:
+        trainer.checkpoint_connector.restore(str(save_path), on_gpu=trainer.on_gpu)
+        fs = get_filesystem(str(save_path))
+        if fs.exists(save_path):
+            fs.rm(save_path)
 
     # Finish by resetting variables so trainer is ready to fit model
     __lr_finder_restore_params(trainer, model)

tests/trainer/flags/test_fast_dev_run.py

@@ -0,0 +1,21 @@
+import pytest
+from pytorch_lightning import Trainer
+from tests.base import EvalModelTemplate
+
+
+@pytest.mark.parametrize('tuner_alg', ['batch size scaler', 'learning rate finder'])
+def test_skip_on_fast_dev_run_batch_scaler(tmpdir, tuner_alg):
+    """ Test that tuner algorithms are skipped if fast dev run is enabled """
+
+    hparams = EvalModelTemplate.get_default_hparams()
+    model = EvalModelTemplate(**hparams)
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=2,
+        auto_scale_batch_size=True if tuner_alg == 'batch size scaler' else False,
+        auto_lr_find=True if tuner_alg == 'learning rate finder' else False,
+        fast_dev_run=True
+    )
+    expected_message = f'Skipping {tuner_alg} since `fast_dev_run=True`'
+    with pytest.warns(UserWarning, match=expected_message):
+        trainer.tune(model)