set_epoch for validation and prediction data loaders (#12197)

* set_epoch for prediction and evaluation
* minor fix in the test, warning msg was changed

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Adrian Wälchli <[email protected]>

Author: 3 people

pytorch_lightning/loops/dataloader/evaluation_loop.py

@@ -159,6 +159,18 @@ def advance(self, *args: Any, **kwargs: Any) -> None:
             # indicate the loop has run
             self._has_run = True
 
+    def on_advance_start(self, *args: Any, **kwargs: Any) -> None:
+        dataloader = self.current_dataloader
+        if (
+            dataloader is not None
+            and getattr(dataloader, "sampler", None)
+            and callable(getattr(dataloader.sampler, "set_epoch", None))
+        ):
+            # set seed for distributed sampler (enables shuffling for each epoch)
+            dataloader.sampler.set_epoch(self.trainer.fit_loop.epoch_progress.current.processed)
+
+        super().on_advance_start(*args, **kwargs)
+
     def on_advance_end(self) -> None:
         self.trainer._logger_connector.epoch_end_reached()
 

pytorch_lightning/loops/dataloader/prediction_loop.py

@@ -87,6 +87,14 @@ def advance(self, *args: Any, **kwargs: Any) -> None:
         """Predicts one entire dataloader."""
         void(*args, **kwargs)
         dataloader = self.current_dataloader
+        if (
+            dataloader is not None
+            and getattr(dataloader, "sampler", None)
+            and callable(getattr(dataloader.sampler, "set_epoch", None))
+        ):
+            # set seed for distributed sampler (enables shuffling for each epoch)
+            dataloader.sampler.set_epoch(self.trainer.fit_loop.epoch_progress.current.processed)
+        dataloader = self.trainer.strategy.process_dataloader(dataloader)
         dataloader_iter = enumerate(dataloader)
         dl_max_batches = self.max_batches[self.current_dataloader_idx]
 

pytorch_lightning/trainer/connectors/data_connector.py

@@ -484,8 +484,8 @@ def replace_sampler(dataloader: DataLoader) -> DataLoader:
     def _check_eval_shuffling(dataloader, mode):
         if _is_dataloader_shuffled(dataloader):
             rank_zero_warn(
-                f"Your `{mode.dataloader_prefix}_dataloader` has `shuffle=True`,"
-                " it is strongly recommended that you turn this off for val/test/predict dataloaders.",
+                f"Your `{mode.dataloader_prefix}_dataloader`'s sampler has shuffling enabled,"
+                " it is strongly recommended that you turn shuffling off for val/test/predict dataloaders.",
                 category=PossibleUserWarning,
             )
 

tests/loops/test_evaluation_loop.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from unittest import mock
+from unittest.mock import Mock
 
 import torch
 from torch.utils.data.dataloader import DataLoader
+from torch.utils.data.sampler import RandomSampler
 
 from pytorch_lightning import Trainer
 from pytorch_lightning.loops import EvaluationEpochLoop
@@ -42,6 +44,34 @@ def test_on_evaluation_epoch_end(eval_epoch_end_mock, tmpdir):
     assert eval_epoch_end_mock.call_count == 4
 
 
+def test_set_epoch_called_eval_predict(tmpdir):
+    """Tests that set_epoch (if the sampler has one) is called on the DataLoader during evaluation and
+    prediction."""
+
+    def _get_dataloader():
+        dataset = RandomDataset(32, 64)
+        sampler = RandomSampler(dataset)
+        sampler.set_epoch = Mock()
+        return DataLoader(dataset, sampler=sampler)
+
+    model = BoringModel()
+    trainer = Trainer(
+        default_root_dir=tmpdir, limit_train_batches=2, limit_val_batches=2, max_epochs=2, enable_model_summary=False
+    )
+
+    train_dataloader = _get_dataloader()
+    val_dataloader = _get_dataloader()
+    trainer.fit(model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader)
+    # One for each epoch
+    assert train_dataloader.sampler.set_epoch.call_count == 2
+    # One for each epoch + sanity check
+    assert val_dataloader.sampler.set_epoch.call_count == 3
+
+    val_dataloader = _get_dataloader()
+    trainer.validate(model, val_dataloader)
+    assert val_dataloader.sampler.set_epoch.call_count == 1
+
+
 @mock.patch(
     "pytorch_lightning.trainer.connectors.logger_connector.logger_connector.LoggerConnector.log_eval_end_metrics"
 )

tests/trainer/test_data_loading.py

@@ -378,5 +378,5 @@ def test_non_sequential_sampler_warning_is_raised_for_eval_dataloader(val_dl):
     trainer = Trainer()
     model = BoringModel()
     trainer._data_connector.attach_data(model, val_dataloaders=val_dl)
-    with pytest.warns(PossibleUserWarning, match="recommended .* turn this off for val/test/predict"):
+    with pytest.warns(PossibleUserWarning, match="recommended .* turn shuffling off for val/test/predict"):
         trainer._data_connector._reset_eval_dataloader(RunningStage.VALIDATING, model)