restore loops and intergrate early exit

Author: rohitgr7

pytorch_lightning/callbacks/batch_size_finder.py

@@ -20,6 +20,7 @@
 
 import os
 import uuid
+from copy import deepcopy
 from typing import Optional, Tuple
 
 from torch.utils.data.dataloader import DataLoader
@@ -38,7 +39,15 @@
 
 
 class BatchSizeFinder(Callback):
-    def __init__(self, mode: str = "power", steps_per_trial=3, init_val=2, max_trials=25, batch_arg_name="batch_size"):
+    def __init__(
+        self,
+        mode: str = "power",
+        steps_per_trial=3,
+        init_val=2,
+        max_trials=25,
+        batch_arg_name="batch_size",
+        early_exit=False,
+    ):
 
         mode = mode.lower()
         if mode not in ("power", "binsearch"):
@@ -50,6 +59,7 @@ def __init__(self, mode: str = "power", steps_per_trial=3, init_val=2, max_trial
         self.max_trials = max_trials
         self.batch_arg_name = batch_arg_name
         self.optimal_batch_size = init_val
+        self.early_exit = early_exit
 
     def scale_batch_size(self, trainer, pl_module):
         if trainer.fast_dev_run:
@@ -90,7 +100,7 @@ def scale_batch_size(self, trainer, pl_module):
         self._reset_params(trainer)
 
         # Save initial model, that is loaded after batch size is found
-        save_path = os.path.join(trainer.default_root_dir, f"scale_batch_size_temp_model_{uuid.uuid4()}.ckpt")
+        save_path = os.path.join(trainer.default_root_dir, f".scale_batch_size_temp_model_{uuid.uuid4()}.ckpt")
         trainer.save_checkpoint(save_path)
 
         if trainer.progress_bar_callback:
@@ -112,6 +122,7 @@ def scale_batch_size(self, trainer, pl_module):
                 fs.rm(save_path)
 
         self._restore_params(trainer)
+
         if trainer.progress_bar_callback:
             trainer.progress_bar_callback.enable()
 
@@ -182,7 +193,12 @@ def _run_binary_scaling(self, trainer, pl_module, new_size):
                     garbage_collection_cuda()
                     high = new_size
                     midval = (high + low) // 2
-                    new_size, _ = self._adjust_batch_size(trainer, value=midval, desc="failed")
+                    new_size, changed = self._adjust_batch_size(trainer, value=midval, desc="failed")
+
+                    if changed:
+                        # Force the dataloaders to reset as the batch size has changed
+                        self._reset_dataloaders(trainer, pl_module)
+
                     if high - low <= 1:
                         break
                 else:
@@ -193,14 +209,17 @@ def _run_binary_scaling(self, trainer, pl_module, new_size):
     def _try_loop_run(self, trainer):
         if trainer.state.fn == TrainerFn.FITTING:
             trainer.fit_loop.global_step = self._dumped_params["global_step"]
-            trainer.fit_loop.current_epoch = self._dumped_params["current_epoch"]
-            trainer.fit_loop.run()
+            # trainer.fit_loop.current_epoch = self._dumped_params["current_epoch"]
+            loop = trainer.fit_loop
         elif trainer.state.fn == TrainerFn.VALIDATING:
-            trainer.validate_loop.run()
+            loop = trainer.validate_loop
         elif trainer.state.fn == TrainerFn.TESTING:
-            trainer.test_loop.run()
+            loop = trainer.test_loop
         elif trainer.state.fn == TrainerFn.PREDICTING:
-            trainer.predict_loop.run()
+            loop = trainer.predict_loop
+
+        loop.load_state_dict(deepcopy(self._dumped_params["loop_state_dict"]))
+        loop.run()
 
     @staticmethod
     def _reset_dataloaders(trainer, pl_module):
@@ -216,57 +235,153 @@ def _reset_dataloaders(trainer, pl_module):
 
     def _dump_params(self, trainer):
         self._dumped_params = {
-            "current_epoch": trainer.current_epoch,
-            "global_step": trainer.global_step,
-            "max_steps": trainer.max_steps,
+            # "current_epoch": trainer.current_epoch,
             "logger": trainer.logger,
             "callbacks": trainer.callbacks,
-            "limit_train_batches": trainer.limit_train_batches,
-            "limit_val_batches": trainer.limit_val_batches,
-            "limit_test_batches": trainer.limit_test_batches,
-            "limit_predict_batches": trainer.limit_predict_batches,
         }
 
+        if trainer.state.fn == TrainerFn.FITTING:
+            loop = trainer.fit_loop
+            self._dumped_params["global_step"] = trainer.global_step
+            self._dumped_params["max_steps"] = trainer.max_steps
+            self._dumped_params["limit_val_batches"] = trainer.limit_val_batches
+        elif trainer.state.fn == TrainerFn.VALIDATING:
+            loop = trainer.validate_loop
+            self._dumped_params["limit_val_batches"] = trainer.limit_val_batches
+        elif trainer.state.fn == TrainerFn.TESTING:
+            loop = trainer.test_loop
+            self._dumped_params["limit_test_batches"] = trainer.limit_test_batches
+        elif trainer.state.fn == TrainerFn.PREDICTING:
+            loop = trainer.predict_loop
+            self._dumped_params["limit_predict_batches"] = trainer.limit_predict_batches
+
+        self._dumped_params["loop_state_dict"] = deepcopy(loop.state_dict())
+        if hasattr(loop, "verbose"):
+            self._dumped_params["loop_verbose"] = loop.verbose
+
     def _reset_params(self, trainer):
         trainer.logger = DummyLogger() if trainer.logger is not None else None
         trainer.callbacks = []
+
         if trainer.state.fn == TrainerFn.FITTING:
             trainer.limit_val_batches = self.steps_per_trial
             trainer.fit_loop.max_steps = self.steps_per_trial
         elif trainer.state.fn == TrainerFn.VALIDATING:
             trainer.limit_val_batches = self.steps_per_trial
+            trainer.validate_loop.verbose = False
         elif trainer.state.fn == TrainerFn.TESTING:
             trainer.limit_test_batches = self.steps_per_trial
+            trainer.test_loop.verbose = False
         elif trainer.state.fn == TrainerFn.PREDICTING:
             trainer.limit_predict_batches = self.steps_per_trial
 
     def _restore_params(self, trainer):
-        trainer.fit_loop.current_epoch = self._dumped_params["current_epoch"]
-        trainer.fit_loop.global_step = self._dumped_params["global_step"]
-        trainer.fit_loop.max_steps = self._dumped_params["max_steps"]
         trainer.logger = self._dumped_params["logger"]
         trainer.callbacks = self._dumped_params["callbacks"]
-        trainer.limit_train_batches = self._dumped_params["limit_train_batches"]
-        trainer.limit_val_batches = self._dumped_params["limit_val_batches"]
-        trainer.limit_test_batches = self._dumped_params["limit_test_batches"]
-        trainer.limit_predict_batches = self._dumped_params["limit_predict_batches"]
 
-    def on_train_epoch_start(self, trainer, pl_module):
-        self.scale_batch_size(trainer, pl_module)
+        if trainer.state.fn == TrainerFn.FITTING:
+            # trainer.fit_loop.current_epoch = self._dumped_params["current_epoch"]
+            trainer.fit_loop.global_step = self._dumped_params["global_step"]
+            loop = trainer.fit_loop
+            loop.max_steps = self._dumped_params["max_steps"]
+            trainer.limit_val_batches = self._dumped_params["limit_val_batches"]
+        elif trainer.state.fn == TrainerFn.VALIDATING:
+            loop = trainer.validate_loop
+            trainer.limit_val_batches = self._dumped_params["limit_val_batches"]
+        elif trainer.state.fn == TrainerFn.TESTING:
+            loop = trainer.test_loop
+            trainer.limit_test_batches = self._dumped_params["limit_test_batches"]
+        elif trainer.state.fn == TrainerFn.PREDICTING:
+            loop = trainer.predict_loop
+            trainer.limit_predict_batches = self._dumped_params["limit_predict_batches"]
+
+        loop.load_state_dict(deepcopy(self._dumped_params["loop_state_dict"]))
+        if "loop_verbose" in self._dumped_params:
+            loop.verbose = self._dumped_params["loop_verbose"]
+
+    def pre_early_exit(self, trainer):
+        if trainer.state.fn == TrainerFn.FITTING:
+            trainer.should_stop = True
+            self._dumped_params["num_training_batches"] = trainer.num_training_batches
+            trainer.num_training_batches = 0
+        elif trainer.state.fn == TrainerFn.VALIDATING:
+            self._dumped_params["num_val_batches"] = trainer.num_val_batches
+            trainer.num_val_batches = [0]
+        elif trainer.state.fn == TrainerFn.TESTING:
+            self._dumped_params["num_test_batches"] = trainer.num_test_batches
+            trainer.num_test_batches = [0]
+        elif trainer.state.fn == TrainerFn.PREDICTING:
+            self._dumped_params["num_predict_batches"] = trainer.num_predict_batches
+            trainer.num_predict_batches = [0]
+
+    def post_early_exit(self, trainer):
+        if trainer.state.fn == TrainerFn.FITTING:
+            trainer.num_training_batches = self._dumped_params["num_training_batches"]
+            loop = trainer.fit_loop
+        if trainer.state.fn == TrainerFn.VALIDATING:
+            trainer.num_val_batches = self._dumped_params["num_val_batches"]
+            loop = trainer.validate_loop
+        if trainer.state.fn == TrainerFn.TESTING:
+            trainer.num_test_batches = self._dumped_params["num_test_batches"]
+            loop = trainer.test_loop
+        if trainer.state.fn == TrainerFn.PREDICTING:
+            trainer.num_predict_batches = self._dumped_params["num_predict_batches"]
+            loop = trainer.predict_loop
+
+        loop.load_state_dict(self._dumped_params["loop_state_dict"])
         trainer.callbacks = [cb for cb in trainer.callbacks if not isinstance(cb, BatchSizeFinder)]
 
-    def on_validation_epoch_start(self, trainer, pl_module):
-        if not trainer.sanity_checking:
-            self.scale_batch_size(trainer, pl_module)
+    def on_fit_start(self, trainer, pl_module):
+        self.scale_batch_size(trainer, pl_module)
+
+        if self.early_exit:
+            self.pre_early_exit(trainer)
+        else:
+            trainer.callbacks = [cb for cb in trainer.callbacks if not isinstance(cb, BatchSizeFinder)]
+
+    def on_validation_start(self, trainer, pl_module):
+        if trainer.sanity_checking or trainer.state.fn != TrainerFn.VALIDATING:
+            return
+
+        if self.early_exit:
+            self.pre_early_exit(trainer)
+        else:
             trainer.callbacks = [cb for cb in trainer.callbacks if not isinstance(cb, BatchSizeFinder)]
 
-    def on_test_epoch_start(self, trainer, pl_module):
+    def on_test_start(self, trainer, pl_module):
         self.scale_batch_size(trainer, pl_module)
-        trainer.callbacks = [cb for cb in trainer.callbacks if not isinstance(cb, BatchSizeFinder)]
 
-    def on_predict_epoch_start(self, trainer, pl_module):
+        if self.early_exit:
+            self.pre_early_exit(trainer)
+        else:
+            trainer.callbacks = [cb for cb in trainer.callbacks if not isinstance(cb, BatchSizeFinder)]
+
+    def on_predict_start(self, trainer, pl_module):
         self.scale_batch_size(trainer, pl_module)
-        trainer.callbacks = [cb for cb in trainer.callbacks if not isinstance(cb, BatchSizeFinder)]
+
+        if self.early_exit:
+            self.pre_early_exit(trainer)
+        else:
+            trainer.callbacks = [cb for cb in trainer.callbacks if not isinstance(cb, BatchSizeFinder)]
+
+    def on_fit_end(self, trainer, pl_module):
+        if self.early_exit:
+            self.post_early_exit(trainer)
+
+    def on_validation_end(self, trainer, pl_module):
+        if trainer.sanity_checking or trainer.state.fn != TrainerFn.VALIDATING:
+            return
+
+        if self.early_exit:
+            self.post_early_exit(trainer)
+
+    def on_test_end(self, trainer, pl_module):
+        if self.early_exit:
+            self.post_early_exit(trainer)
+
+    def on_predict_end(self, trainer, pl_module):
+        if self.early_exit:
+            self.post_early_exit(trainer)
 
     def _adjust_batch_size(
         self,
@@ -295,19 +410,19 @@ def _adjust_batch_size(
         if desc:
             rank_zero_info(f"Batch size {batch_size} {desc}, trying batch size {new_size}")
 
-        # TODO improve this for CombinedLoader
+        # TODO improve this for CombinedLoader and multi dataloaders
         if trainer.state.fn == TrainerFn.FITTING:
             if not self._is_valid_batch_size(new_size, trainer.train_dataloader, trainer):
                 new_size = min(new_size, len(trainer.train_dataloader.dataset))
         if trainer.state.fn == TrainerFn.VALIDATING:
             if not self._is_valid_batch_size(new_size, trainer.val_dataloaders, trainer):
-                new_size = min(new_size, len(trainer.val_dataloaders.dataset))
+                new_size = min(new_size, len(trainer.val_dataloaders[0].dataset))
         if trainer.state.fn == TrainerFn.TESTING:
             if not self._is_valid_batch_size(new_size, trainer.test_dataloaders, trainer):
-                new_size = min(new_size, len(trainer.test_dataloaders.dataset))
+                new_size = min(new_size, len(trainer.test_dataloaders[0].dataset))
         if trainer.state.fn == TrainerFn.PREDICTING:
             if not self._is_valid_batch_size(new_size, trainer.predict_dataloaders, trainer):
-                new_size = min(new_size, len(trainer.predict_dataloaders.dataset))
+                new_size = min(new_size, len(trainer.predict_dataloaders[0].dataset))
 
         changed = new_size != batch_size
         lightning_setattr(model, self.batch_arg_name, new_size)
@@ -316,4 +431,4 @@ def _adjust_batch_size(
     @staticmethod
     def _is_valid_batch_size(batch_size: int, dataloader: DataLoader, trainer: "pl.Trainer"):
         module = trainer.lightning_module or trainer.datamodule
-        return not has_len_all_ranks(dataloader, trainer.training_type_plugin, module) or batch_size <= len(dataloader)
+        return not has_len_all_ranks(dataloader, trainer.strategy, module) or batch_size <= len(dataloader)

pytorch_lightning/loops/base.py

@@ -23,6 +23,7 @@
 from pytorch_lightning.trainer.progress import BaseProgress
 from pytorch_lightning.utilities.enums import _FaultTolerantMode
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.imports import _fault_tolerant_training
 
 T = TypeVar("T")  # the output type of `run`
 
@@ -335,4 +336,5 @@ def _load_from_state_dict(self, state_dict: Dict, prefix: str, metrics: Optional
                     v.reset(metrics=False)
 
         self.on_load_checkpoint(state_dict[prefix + "state_dict"])
-        self.restarting = True
+        if _fault_tolerant_training():
+            self.restarting = True

pytorch_lightning/tuner/tuning.py

@@ -204,7 +204,7 @@ def lr_find(
 
     def fit(self, model, train_dataloaders, val_dataloaders, datamodule, **batch_size_scale_kwargs):
         self.trainer.state.fn = None
-        batch_size_finder = BatchSizeFinder(**batch_size_scale_kwargs)
+        batch_size_finder = BatchSizeFinder(**batch_size_scale_kwargs, early_exit=True)
         self.trainer.callbacks = [batch_size_finder] + self.trainer.callbacks
         self.trainer.fit(model, train_dataloaders, val_dataloaders, datamodule)
         return batch_size_finder.optimal_batch_size

tests/tuner/test_scale_batch_size.py

@@ -92,7 +92,7 @@ def test_model_reset_correctly(tmpdir):
             torch.eq(before_state_dict[key], after_state_dict[key])
         ), "Model was not reset correctly after scaling batch size"
 
-    assert not any(f for f in os.listdir(tmpdir) if f.startswith("scale_batch_size_temp_model"))
+    assert not any(f for f in os.listdir(tmpdir) if f.startswith(".scale_batch_size_temp_model"))
 
 
 def test_trainer_reset_correctly(tmpdir):
@@ -105,18 +105,19 @@ def test_trainer_reset_correctly(tmpdir):
     trainer = Trainer(default_root_dir=tmpdir, max_epochs=1)
 
     changed_attributes = [
-        "callbacks",
-        "checkpoint_callback",
-        "current_epoch",
-        "limit_train_batches",
-        "logger",
-        "max_steps",
         "global_step",
+        "limit_val_batches",
+        "max_steps",
+        "logger",
+        "callbacks",
     ]
     expected = {ca: getattr(trainer, ca) for ca in changed_attributes}
-    trainer.tuner.scale_batch_size(model, max_trials=5)
+    expected_loop_state_dict = trainer.fit_loop.state_dict()
+    trainer.tuner.scale_batch_size(model, max_trials=64)
     actual = {ca: getattr(trainer, ca) for ca in changed_attributes}
+    actual_loop_state_dict = trainer.fit_loop.state_dict()
 
+    assert expected_loop_state_dict == actual_loop_state_dict
     assert actual == expected