fix mypy

Author: rohitgr7

pytorch_lightning/callbacks/batch_size_finder.py

@@ -21,7 +21,7 @@
 import os
 import uuid
 from copy import deepcopy
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple, TypedDict, Union
 
 from torch.utils.data.dataloader import DataLoader
 
@@ -87,6 +87,18 @@ def __init__(
 
         self._early_exit = False
 
+        from pytorch_lightning.loggers.base import LightningLoggerBase
+
+        class _BatchSizeFinderDumpedParams(TypedDict):
+            callbacks: List[Callback]
+            logger: Optional[LightningLoggerBase]
+            max_steps: int
+            global_step: Optional[None]
+            limit_val_batches: Union[int, float]
+            limit_eval_batches: Union[int, float]
+
+        self._dumped_params: _BatchSizeFinderDumpedParams = {}
+
     def scale_batch_size(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
         if trainer.fast_dev_run:
             rank_zero_warn("Skipping batch size scaler since `fast_dev_run` is enabled.")
@@ -284,13 +296,13 @@ def _dump_params(self, trainer: "pl.Trainer") -> None:
             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
+            self._dumped_params["limit_eval_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
+            self._dumped_params["limit_eval_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["limit_eval_batches"] = trainer.limit_predict_batches
 
         self._dumped_params["loop_state_dict"] = deepcopy(loop.state_dict(force_save_progress=True))
         if hasattr(loop, "verbose"):
@@ -328,13 +340,13 @@ def _restore_params(self, trainer: "pl.Trainer") -> None:
             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"]
+            trainer.limit_val_batches = self._dumped_params["limit_eval_batches"]
         elif trainer.state.fn == TrainerFn.TESTING:
             loop = trainer.test_loop
-            trainer.limit_test_batches = self._dumped_params["limit_test_batches"]
+            trainer.limit_test_batches = self._dumped_params["limit_eval_batches"]
         elif trainer.state.fn == TrainerFn.PREDICTING:
             loop = trainer.predict_loop
-            trainer.limit_predict_batches = self._dumped_params["limit_predict_batches"]
+            trainer.limit_predict_batches = self._dumped_params["limit_eval_batches"]
 
         loop.load_state_dict(deepcopy(self._dumped_params["loop_state_dict"]))
         if "loop_verbose" in self._dumped_params:
@@ -386,18 +398,18 @@ 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 and multi dataloaders
+        # TODO improve this for eval 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):
+            if not self._is_valid_batch_size(new_size, trainer.val_dataloaders[0], trainer):
                 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):
+            if not self._is_valid_batch_size(new_size, trainer.test_dataloaders[0], trainer):
                 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):
+            if not self._is_valid_batch_size(new_size, trainer.predict_dataloaders[0], trainer):
                 new_size = min(new_size, len(trainer.predict_dataloaders[0].dataset))
 
         changed = new_size != batch_size

setup.cfg

@@ -24,13 +24,13 @@ addopts =
     --doctest-modules
     --color=yes
     --disable-pytest-warnings
-filterwarnings =
+# filterwarnings =
     # error out on our deprecation warnings - ensures the code and tests are kept up-to-date
-    error::pytorch_lightning.utilities.warnings.LightningDeprecationWarning
+#     error::pytorch_lightning.utilities.warnings.LightningDeprecationWarning
     # warnings from deprecated modules on import
     # TODO: remove in 1.7
-    ignore::pytorch_lightning.utilities.warnings.LightningDeprecationWarning:pytorch_lightning.core.decorators
-    ignore::pytorch_lightning.utilities.warnings.LightningDeprecationWarning:pytorch_lightning.core.memory
+#     ignore::pytorch_lightning.utilities.warnings.LightningDeprecationWarning:pytorch_lightning.core.decorators
+#     ignore::pytorch_lightning.utilities.warnings.LightningDeprecationWarning:pytorch_lightning.core.memory
 
 junit_duration_report = call
 

tests/tuner/test_scale_batch_size.py

@@ -79,54 +79,45 @@ def test_scale_batch_size_method_with_model_or_datamodule(tmpdir, model_bs, dm_b
         assert trainer.train_dataloader.loaders.batch_size == new_batch_size
 
 
-def test_model_reset_correctly(tmpdir):
-    """Check that model weights are correctly reset after scaling batch size."""
+@pytest.mark.parametrize("trainer_fn", ["fit", "validate", "test", "predict"])
+def test_trainer_reset_correctly(tmpdir, trainer_fn):
+    """Check that model and all trainer parameters are reset correctly after scaling batch size."""
     tutils.reset_seed()
 
     model = BatchSizeModel(batch_size=2)
-
-    # logger file to get meta
-    trainer = Trainer(default_root_dir=tmpdir, max_epochs=1)
-
     before_state_dict = deepcopy(model.state_dict())
 
-    trainer.tuner.scale_batch_size(model, max_trials=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 scaling batch size"
-
-    assert not any(f for f in os.listdir(tmpdir) if f.startswith(".scale_batch_size_temp_model"))
-
-
-def test_trainer_reset_correctly(tmpdir):
-    """Check that all trainer parameters are reset correctly after scaling batch size."""
-    tutils.reset_seed()
-
-    model = BatchSizeModel(batch_size=2)
-
     # logger file to get meta
     trainer = Trainer(default_root_dir=tmpdir, max_epochs=1)
 
     changed_attributes = [
+        "logger",
+        "callbacks",
         "global_step",
         "limit_val_batches",
         "max_steps",
-        "logger",
-        "callbacks",
+        "limit_val_batches",
+        "limit_test_batches",
+        "limit_predict_batches",
     ]
+
     expected = {ca: getattr(trainer, ca) for ca in changed_attributes}
     expected_loop_state_dict = trainer.fit_loop.state_dict()
-    trainer.tuner.scale_batch_size(model, max_trials=64)
+    trainer.tuner.scale_batch_size(model, max_trials=64, method=trainer_fn)
     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
 
+    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 scaling batch size"
+
+    assert not any(f for f in os.listdir(tmpdir) if f.startswith(".scale_batch_size_temp_model"))
+
 
 @RunIf(min_gpus=1)
 @pytest.mark.parametrize("scale_arg", ["power", "binsearch", True])