Moved changes to TrainerDataLoadingMixing

Author: adamviola

pytorch_lightning/trainer/data_loading.py

@@ -47,6 +47,7 @@ class TrainerDataLoadingMixin(ABC):
     # this is just a summary on variables used in this abstract class,
     #  the proper values/initialisation should be done in child class
     val_check_interval: float
+    reload_dataloaders_every_n_epochs: int
     tpu_local_core_rank: int
     train_dataloader: DataLoader
     limit_train_batches: Union[int, float]
@@ -67,6 +68,9 @@ class TrainerDataLoadingMixin(ABC):
     accelerator: Accelerator
     call_hook: Callable
     _accelerator_connector: AcceleratorConnector
+    current_epoch: int
+    _last_train_dl_reload_epoch: int
+    _last_val_dl_reload_epoch: int
 
     def _worker_check(self, dataloader: DataLoader, name: str) -> None:
         if not isinstance(dataloader, DataLoader):
@@ -277,6 +281,9 @@ def reset_train_dataloader(self, model: Optional["pl.LightningModule"] = None) -
                 category=PossibleUserWarning,
             )
 
+        # store epoch of dataloader reset for reload_dataloaders_every_n_epochs
+        self._last_train_dl_reload_epoch = self.current_epoch
+
     def _reset_eval_dataloader(
         self, mode: RunningStage, model: Optional["pl.LightningModule"] = None
     ) -> Tuple[List[Union[int, float]], List[DataLoader]]:
@@ -372,6 +379,9 @@ def reset_val_dataloader(self, model: Optional["pl.LightningModule"] = None) ->
                 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.
 
@@ -459,3 +469,15 @@ def replace_sampler(dataloader: DataLoader) -> DataLoader:
             dataloader = apply_to_collection(dataloader, DataLoader, replace_sampler)
 
         return dataloader
+
+    @property
+    def _should_reload_train_dl(self) -> bool:
+        """Check if train dataloader should be reloaded."""
+        n_epochs = self.reload_dataloaders_every_n_epochs
+        return n_epochs and self.current_epoch - self._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.reload_dataloaders_every_n_epochs
+        return n_epochs and self.current_epoch - self._last_val_dl_reload_epoch >= n_epochs

pytorch_lightning/trainer/trainer.py

@@ -1358,14 +1358,6 @@ def _run_sanity_check(self, ref_model):
             # restore the previous stage when the sanity check if finished
             self.state.stage = stage
 
-    def reset_train_dataloader(self, model: Optional["pl.LightningModule"] = None) -> None:
-        self._last_train_dl_reload_epoch = self.current_epoch
-        return super().reset_train_dataloader(model=model)
-
-    def reset_val_dataloader(self, model: Optional["pl.LightningModule"] = None) -> None:
-        self._last_val_dl_reload_epoch = self.current_epoch
-        return super().reset_val_dataloader(model=model)
-
     def __set_ckpt_path(self, ckpt_path: Optional[str], model_provided: bool, model_connected: bool) -> Optional[str]:
         if model_provided and ckpt_path is None:
             # use passed model to function without loading weights
@@ -1850,18 +1842,6 @@ def progress_bar_dict(self) -> dict:
             return self.progress_bar_callback.get_metrics(self, ref_model)
         return self.progress_bar_metrics
 
-    @property
-    def _should_reload_train_dl(self) -> bool:
-        """Check if train dataloader should be reloaded."""
-        n_epochs = self.reload_dataloaders_every_n_epochs
-        return n_epochs and self.current_epoch - self._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.reload_dataloaders_every_n_epochs
-        return n_epochs and self.current_epoch - self._last_val_dl_reload_epoch >= n_epochs
-
     @property
     def enable_validation(self) -> bool:
         """Check if we should run validation during training."""