Fix an issue to avoid the impact of sanity check on reload_dataloaders_every_n_epochs for validation (#13964)

Author: rohitgr7

src/pytorch_lightning/CHANGELOG.md

@@ -43,6 +43,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed not preserving set attributes on `DataLoader` and `BatchSampler` when instantiated inside `*_dataloader` hooks ([#14212](https://github.com/Lightning-AI/lightning/pull/14212))
 
 
+- Fixed an issue to avoid the impact of sanity check on `reload_dataloaders_every_n_epochs` for validation ([#13964](https://github.com/Lightning-AI/lightning/pull/13964))
+
+
 ## [1.7.1] - 2022-08-09
 
 ### Fixed

src/pytorch_lightning/loops/dataloader/prediction_loop.py

@@ -60,7 +60,8 @@ def max_batches(self) -> List[int]:
     @property
     def dataloaders(self) -> Sequence[DataLoader]:
         """Returns all prediction dataloaders."""
-        return self.trainer.predict_dataloaders
+        dataloaders = self.trainer.predict_dataloaders
+        return [] if dataloaders is None else dataloaders
 
     @property
     def skip(self) -> bool:

src/pytorch_lightning/loops/epoch/prediction_epoch_loop.py

@@ -162,8 +162,9 @@ def _get_batch_indices(self, dataloader_idx: int) -> List[List[int]]:
         """Returns a reference to the seen batch indices if the dataloader has a batch sampler wrapped by our
         :class:`~pytorch_lightning.overrides.distributed.IndexBatchSamplerWrapper`."""
         # the batch_sampler is not be defined in case of CombinedDataLoaders
+        assert self.trainer.predict_dataloaders
         batch_sampler = getattr(
-            self.trainer.predict_dataloaders[dataloader_idx],  # type: ignore[has-type]
+            self.trainer.predict_dataloaders[dataloader_idx],
             "batch_sampler",
             None,
         )

src/pytorch_lightning/loops/fit_loop.py

@@ -210,7 +210,8 @@ def on_run_start(self) -> None:  # type: ignore[override]
 
         self.trainer.reset_train_dataloader(self.trainer.lightning_module)
         # reload the evaluation dataloaders too for proper display in the progress bar
-        self.epoch_loop.val_loop._reload_evaluation_dataloaders()
+        if self.epoch_loop._should_check_val_epoch():
+            self.epoch_loop.val_loop._reload_evaluation_dataloaders()
 
         data_fetcher_cls = _select_data_fetcher(self.trainer)
         self._data_fetcher = data_fetcher_cls(prefetch_batches=self.prefetch_batches)

src/pytorch_lightning/trainer/connectors/data_connector.py

@@ -61,13 +61,19 @@ def __init__(self, trainer: "pl.Trainer", multiple_trainloader_mode: str = "max_
     def _should_reload_train_dl(self) -> bool:
         """Check if train dataloader should be reloaded."""
         n_epochs = self.trainer.reload_dataloaders_every_n_epochs
-        return n_epochs and (self.trainer.current_epoch - self.trainer._last_train_dl_reload_epoch >= n_epochs)
+        return n_epochs and (
+            self.trainer._last_train_dl_reload_epoch is None
+            or self.trainer.current_epoch - self.trainer._last_train_dl_reload_epoch >= n_epochs
+        )
 
     @property
     def _should_reload_val_dl(self) -> bool:
         """Check if validation dataloader should be reloaded."""
         n_epochs = self.trainer.reload_dataloaders_every_n_epochs
-        return n_epochs and (self.trainer.current_epoch - self.trainer._last_val_dl_reload_epoch >= n_epochs)
+        return n_epochs and (
+            self.trainer._last_val_dl_reload_epoch is None
+            or self.trainer.current_epoch - self.trainer._last_val_dl_reload_epoch >= n_epochs
+        )
 
     def on_trainer_init(
         self,

src/pytorch_lightning/trainer/trainer.py

@@ -626,12 +626,12 @@ def _setup_on_init(self) -> None:
         self.num_sanity_val_batches = []
         self.num_test_batches = []
         self.num_val_batches = []
+        self.num_predict_batches = []
         self.test_dataloaders = None
         self.val_dataloaders = None
-        self._last_train_dl_reload_epoch = float("-inf")
-        self._last_val_dl_reload_epoch = float("-inf")
-
-        self.num_predict_batches = []
+        self.predict_dataloaders = None
+        self._last_train_dl_reload_epoch = None
+        self._last_val_dl_reload_epoch: Optional[int] = None
 
     def _call_and_handle_interrupt(self, trainer_fn: Callable, *args: Any, **kwargs: Any) -> Any:
         r"""
@@ -711,8 +711,6 @@ def _fit_impl(
         self.state.fn = TrainerFn.FITTING
         self.state.status = TrainerStatus.RUNNING
         self.training = True
-        self._last_train_dl_reload_epoch = float("-inf")
-        self._last_val_dl_reload_epoch = float("-inf")
 
         # if a datamodule comes in as the second arg, then fix it for the user
         if isinstance(train_dataloaders, LightningDataModule):
@@ -1923,13 +1921,18 @@ def reset_val_dataloader(self, model: Optional["pl.LightningModule"] = None) ->
         has_step = is_overridden("validation_step", pl_module)
         enable_validation = self.limit_val_batches > 0
         if source.is_defined() and has_step and enable_validation:
+            # store epoch of dataloader reset for reload_dataloaders_every_n_epochs
+            # it should not reload again if it has already reloaded during sanity_check
+            if self.state.fn == TrainerFn.FITTING and (
+                (self.sanity_checking and self.fit_loop.epoch_loop._should_check_val_epoch())
+                or not self.sanity_checking
+            ):
+                self._last_val_dl_reload_epoch = self.current_epoch
+
             self.num_val_batches, self.val_dataloaders = self._data_connector._reset_eval_dataloader(
                 RunningStage.VALIDATING, model=pl_module
             )
 
-            # store epoch of dataloader reset for reload_dataloaders_every_n_epochs
-            self._last_val_dl_reload_epoch = self.current_epoch
-
     def reset_test_dataloader(self, model: Optional["pl.LightningModule"] = None) -> None:
         """Resets the test dataloader and determines the number of batches.
 

tests/tests_pytorch/trainer/test_dataloaders.py

@@ -946,104 +946,87 @@ def test_dataloaders_load_only_once_no_sanity_check(tmpdir):
     assert tracker.mock_calls == expected_sequence
 
 
-@pytest.mark.parametrize("n", [1, 2])
-def test_dataloaders_load_every_n_epochs(tmpdir, n):
-    train_reload_epochs, val_reload_epochs = [], []
-
-    class TestModel(BoringModel):
-        def train_dataloader(self):
-            train_reload_epochs.append(self.current_epoch)
-            return super().train_dataloader()
-
-        def val_dataloader(self):
-            val_reload_epochs.append(self.current_epoch)
-            return super().val_dataloader()
-
-    model = TestModel()
-
-    trainer = Trainer(
-        default_root_dir=tmpdir,
-        limit_train_batches=0.3,
-        limit_val_batches=0.3,
-        reload_dataloaders_every_n_epochs=n,
-        max_epochs=5,
-    )
-
-    tracker = Mock()
-    model.train_dataloader = Mock(wraps=model.train_dataloader)
-    model.val_dataloader = Mock(wraps=model.val_dataloader)
-    model.test_dataloader = Mock(wraps=model.test_dataloader)
-
-    tracker.attach_mock(model.train_dataloader, "train_dataloader")
-    tracker.attach_mock(model.val_dataloader, "val_dataloader")
-    tracker.attach_mock(model.test_dataloader, "test_dataloader")
-
-    trainer.fit(model)
-    trainer.test(model)
-
-    # Verify the sequence
-    expected_sequence = [call.val_dataloader(), call.train_dataloader()]  # Sanity check first
-    if n == 1:
-        expected_sequence += [call.train_dataloader(), call.val_dataloader()] * 4
-    elif n == 2:
-        expected_sequence += [call.train_dataloader(), call.val_dataloader()] * 2
-    expected_sequence += [call.test_dataloader()]
-
-    assert tracker.mock_calls == expected_sequence
-
-    # Verify epoch of reloads
-    if n == 1:
-        assert train_reload_epochs == [0, 1, 2, 3, 4]
-        assert val_reload_epochs == [0, 1, 2, 3, 4]
-    elif n == 2:
-        assert train_reload_epochs == [0, 2, 4]
-        assert val_reload_epochs == [0, 2, 4]
-
-
 @pytest.mark.parametrize(
-    "n, train_reload_epochs_expect, val_reload_epochs_expect",
+    (
+        "num_sanity_val_steps, check_val_every_n_epoch, reload_dataloaders_every_n_epochs,"
+        " train_reload_epochs_expect,val_reload_epochs_expect,val_step_epochs_expect"
+    ),
     [
-        # Sanity check at epoch 0 creates a validation dataloader, but validation is
-        # checked (and in this case reloaded) every n epochs starting from epoch n-1
-        (3, [0, 2, 4, 6, 8], [0, 2, 5, 8]),
-        (5, [0, 2, 4, 6, 8], [0, 4, 9]),
+        # general case where sanity check reloads the dataloaders for validation on current_epoch=0
+        (0, 1, 1, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
+        (1, 1, 1, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
+        # case where check_val_every_n_epoch < reload_dataloaders_every_n_epochs so expected val_reload_epoch
+        # and val_step_epoch will be different
+        (0, 1, 2, [0, 2, 4, 6, 8], [0, 2, 4, 6, 8], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
+        (1, 1, 2, [0, 2, 4, 6, 8], [2, 4, 6, 8], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
+        (0, 3, 4, [0, 4, 8], [2, 8], [2, 5, 8]),
+        (1, 3, 4, [0, 4, 8], [2, 8], [2, 5, 8]),
+        # case where check_val_every_n_epoch > reload_dataloaders_every_n_epochs so expected val_reload_epoch
+        # and val_step_epoch will be same
+        (0, 2, 1, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [1, 3, 5, 7, 9], [1, 3, 5, 7, 9]),
+        (1, 2, 1, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [1, 3, 5, 7, 9], [1, 3, 5, 7, 9]),
+        (0, 3, 2, [0, 2, 4, 6, 8], [2, 5, 8], [2, 5, 8]),
+        (1, 3, 2, [0, 2, 4, 6, 8], [2, 5, 8], [2, 5, 8]),
+        (0, 5, 2, [0, 2, 4, 6, 8], [4, 9], [4, 9]),
+        (1, 5, 2, [0, 2, 4, 6, 8], [4, 9], [4, 9]),
+        # case where check_val_every_n_epoch = reload_dataloaders_every_n_epochs so expected val_reload_epoch
+        # and val_step_epoch will be same
+        (0, 2, 2, [0, 2, 4, 6, 8], [1, 3, 5, 7, 9], [1, 3, 5, 7, 9]),
+        (1, 2, 2, [0, 2, 4, 6, 8], [1, 3, 5, 7, 9], [1, 3, 5, 7, 9]),
     ],
 )
 def test_dataloaders_load_every_n_epochs_infrequent_val(
-    tmpdir, n, train_reload_epochs_expect, val_reload_epochs_expect
+    tmpdir,
+    num_sanity_val_steps,
+    check_val_every_n_epoch,
+    reload_dataloaders_every_n_epochs,
+    train_reload_epochs_expect,
+    val_reload_epochs_expect,
+    val_step_epochs_expect,
 ):
     """Test dataloader reload behavior when infrequently checking validation set (via check_val_every_n_epoch)"""
-    train_reload_epochs, val_reload_epochs = [], []
+    sanity_val_check_epochs, train_reload_epochs, val_reload_epochs = [], [], []
+    sanity_val_step_epochs, val_step_epochs = [], []
 
     class TestModel(BoringModel):
         def train_dataloader(self):
             train_reload_epochs.append(self.current_epoch)
             return super().train_dataloader()
 
         def val_dataloader(self):
-            val_reload_epochs.append(self.current_epoch)
+            if self.trainer.sanity_checking:
+                sanity_val_check_epochs.append(self.current_epoch)
+            else:
+                val_reload_epochs.append(self.current_epoch)
             return super().val_dataloader()
 
+        def validation_step(self, *args, **kwargs):
+            if self.trainer.sanity_checking:
+                sanity_val_step_epochs.append(self.current_epoch)
+            else:
+                val_step_epochs.append(self.current_epoch)
+
+            return super().validation_step(*args, **kwargs)
+
     model = TestModel()
 
     trainer = Trainer(
         default_root_dir=tmpdir,
-        limit_train_batches=0.3,
-        limit_val_batches=0.3,
-        check_val_every_n_epoch=n,
-        reload_dataloaders_every_n_epochs=2,
+        limit_train_batches=1,
+        limit_val_batches=1,
+        check_val_every_n_epoch=check_val_every_n_epoch,
+        reload_dataloaders_every_n_epochs=reload_dataloaders_every_n_epochs,
         max_epochs=10,
+        num_sanity_val_steps=num_sanity_val_steps,
     )
-    model.test_dataloader = Mock(wraps=model.test_dataloader)
-
     trainer.fit(model)
-    trainer.test(model)
 
     # Verify epoch of reloads
+    sanity_val_check_epochs_expect = [0] if num_sanity_val_steps else []
+    assert sanity_val_check_epochs == sanity_val_step_epochs == sanity_val_check_epochs_expect
     assert train_reload_epochs == train_reload_epochs_expect
     assert val_reload_epochs == val_reload_epochs_expect
-
-    model.test_dataloader.assert_called_once()
+    assert val_step_epochs == val_step_epochs_expect
 
 
 def test_dataloaders_load_every_n_epochs_frequent_val(tmpdir):