Fix lr_find to generate same results on multiple calls (#9704)

Author: rohitgr7

CHANGELOG.md

@@ -11,6 +11,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed gradient unscaling being called too late, causing gradient clipping and gradient norm tracking to be applied incorrectly ([#9606](https://github.com/PyTorchLightning/pytorch-lightning/pull/9606))
 
 
+- Fixed `lr_find` to generate same results on multiple calls ([#9704](https://github.com/PyTorchLightning/pytorch-lightning/pull/9704))
+
+
 ## [1.4.8] - 2021-09-22
 
 - Fixed error reporting in DDP process reconciliation when processes are launched by an external agent (#9389)

pytorch_lightning/tuner/batch_size_scaling.py

@@ -101,6 +101,7 @@ def __scale_batch_dump_params(trainer: "pl.Trainer") -> None:
     trainer.__dumped_params = {
         "auto_lr_find": trainer.auto_lr_find,
         "current_epoch": trainer.current_epoch,
+        "global_step": trainer.global_step,
         "max_steps": trainer.max_steps,
         "weights_summary": trainer.weights_summary,
         "logger": trainer.logger,
@@ -128,6 +129,7 @@ def __scale_batch_reset_params(trainer: "pl.Trainer", model: "pl.LightningModule
 def __scale_batch_restore_params(trainer: "pl.Trainer") -> None:
     trainer.auto_lr_find = trainer.__dumped_params["auto_lr_find"]
     trainer.fit_loop.current_epoch = trainer.__dumped_params["current_epoch"]
+    trainer.fit_loop.global_step = trainer.__dumped_params["global_step"]
     trainer.fit_loop.max_steps = trainer.__dumped_params["max_steps"]
     trainer.weights_summary = trainer.__dumped_params["weights_summary"]
     trainer.logger = trainer.__dumped_params["logger"]

pytorch_lightning/tuner/lr_finder.py

@@ -284,6 +284,7 @@ def __lr_finder_dump_params(trainer, model):
         "auto_lr_find": trainer.auto_lr_find,
         "callbacks": trainer.callbacks,
         "logger": trainer.logger,
+        "global_step": trainer.global_step,
         "max_steps": trainer.max_steps,
         "checkpoint_callback": trainer.checkpoint_callback,
         "current_epoch": trainer.current_epoch,
@@ -295,6 +296,7 @@ def __lr_finder_restore_params(trainer, model):
     trainer.auto_lr_find = trainer.__dumped_params["auto_lr_find"]
     trainer.logger = trainer.__dumped_params["logger"]
     trainer.callbacks = trainer.__dumped_params["callbacks"]
+    trainer.fit_loop.global_step = trainer.__dumped_params["global_step"]
     trainer.fit_loop.max_steps = trainer.__dumped_params["max_steps"]
     trainer.fit_loop.current_epoch = trainer.__dumped_params["current_epoch"]
     model.configure_optimizers = trainer.__dumped_params["configure_optimizers"]

tests/tuner/test_lr_finder.py

@@ -75,6 +75,7 @@ def test_trainer_reset_correctly(tmpdir):
         "checkpoint_callback",
         "current_epoch",
         "logger",
+        "global_step",
         "max_steps",
     ]
     expected = {ca: getattr(trainer, ca) for ca in changed_attributes}
@@ -282,3 +283,17 @@ def training_step_end(self, outputs):
     trainer = Trainer(default_root_dir=tmpdir)
     num_training = 3
     trainer.tuner.lr_find(model=model, num_training=num_training)
+
+
+def test_multiple_lr_find_calls_gives_same_results(tmpdir):
+    """Tests that lr_finder gives same results if called multiple times."""
+    model = BoringModel()
+
+    trainer = Trainer(default_root_dir=tmpdir, max_epochs=2)
+    all_res = [trainer.tuner.lr_find(model).results for _ in range(3)]
+
+    assert all(
+        all_res[0][k] == curr_lr_finder[k] and len(curr_lr_finder[k]) > 10
+        for curr_lr_finder in all_res[1:]
+        for k in all_res[0].keys()
+    )

tests/tuner/test_scale_batch_size.py

@@ -108,6 +108,7 @@ def test_trainer_reset_correctly(tmpdir):
         "limit_train_batches",
         "logger",
         "max_steps",
+        "global_step",
         "weights_summary",
     ]
     expected = {ca: getattr(trainer, ca) for ca in changed_attributes}