Loop specialization

pytorch_lightning/loops/epoch/training_epoch_loop.py

@@ -77,13 +77,12 @@ def done(self) -> bool:
         max_steps_reached = self.max_steps is not None and self.global_step >= self.max_steps
         return max_steps_reached or self.trainer.should_stop or self._num_training_batches_reached(self.is_last_batch)
 
-    def connect(self, trainer: 'pl.Trainer', *args: Any, **kwargs: Any) -> None:
+    def connect(self, trainer: 'pl.Trainer', batch_loop, val_loop) -> None:
         """Connects the loop with all necessary parts like trainer and accelerators"""
-        super().connect(trainer, *args, **kwargs)
-        self.batch_loop = TrainingBatchLoop()
-        self.batch_loop.connect(trainer)
-        self.val_loop = loops.EvaluationLoop()
-        self.val_loop.connect(trainer)
+        super().connect(trainer)
+        self.batch_loop = batch_loop  #  or TrainingBatchLoop()
+        self.val_loop = val_loop  #or loops.EvaluationLoop()
+        # self.val_loop.connect(trainer)
 
     def reset(self) -> None:
         """Resets the internal state of the loop for a new run"""

pytorch_lightning/loops/fit_loop.py

@@ -32,25 +32,18 @@ class FitLoop(Loop):
     Args:
         min_epochs: The minimum number of epochs
         max_epochs: The maximum number of epochs
-        min_steps: The minimum number of steps
-        max_steps: The maximum number of epoch
 
     .. note::
         If neither the minimum epochs nor steps are specified the minimum number of epochs is set to 1
         and if neither the maximum steps nor epochs are specified, the maximum epochs are set to 1000.
     """
 
-    def __init__(
-        self,
-        min_epochs: Optional[int] = None,
-        max_epochs: Optional[int] = None,
-        min_steps: Optional[int] = None,
-        max_steps: Optional[int] = None
-    ):
+    # FIXME: update the note above
+    def __init__(self, min_epochs: Optional[int] = None, max_epochs: Optional[int] = None):
         super().__init__()
-        self.max_epochs = 1000 if (max_epochs is None and max_steps is None) else max_epochs
-        self.min_epochs = 1 if (min_epochs is None and min_steps is None) else min_epochs
-        self.epoch_loop = TrainingEpochLoop(min_steps, max_steps)
+        self.max_epochs = min_epochs
+        self.min_epochs = max_epochs
+        self.epoch_loop = None
 
     @property
     def results(self) -> ResultCollection:
@@ -156,10 +149,10 @@ def skip(self) -> bool:
         """Whether we should skip the training and immediately return from the call to :meth:`run`."""
         return self.done or self.trainer.num_training_batches == 0
 
-    def connect(self, trainer: 'pl.Trainer', *args: Any, **kwargs: Any) -> None:
+    def connect(self, trainer: 'pl.Trainer', epoch_loop) -> None:
         """Connects the loop with necessary arguments like the trainer"""
-        super().connect(trainer, *args, **kwargs)
-        self.epoch_loop.connect(trainer)
+        super().connect(trainer)
+        self.epoch_loop = epoch_loop
 
     def reset(self) -> None:
         """Resets the internal state of this loop"""

pytorch_lightning/trainer/trainer.py

@@ -28,7 +28,10 @@
 from pytorch_lightning.core.datamodule import LightningDataModule
 from pytorch_lightning.core.memory import ModelSummary
 from pytorch_lightning.loggers import LightningLoggerBase
-from pytorch_lightning.loops import EvaluationLoop, FitLoop, PredictionLoop
+from pytorch_lightning.loops import TrainingBatchLoop, TrainingEpochLoop
+from pytorch_lightning.loops.dataloader.evaluation_loop import EvaluationLoop
+from pytorch_lightning.loops.dataloader.prediction_loop import PredictionLoop
+from pytorch_lightning.loops.fit_loop import FitLoop
 from pytorch_lightning.plugins import Plugin
 from pytorch_lightning.plugins.environments import ClusterEnvironment
 from pytorch_lightning.profiler import (
@@ -343,14 +346,30 @@ def __init__(
         self.slurm_connector = SLURMConnector(self)
         self.tuner = Tuner(self)
 
-        self.fit_loop = FitLoop(min_epochs, max_epochs, min_steps, max_steps)
+        # .fit() loop
+        self.fit_loop = FitLoop(
+            min_epochs=(1 if (min_epochs is None and min_steps is None) else min_epochs),
+            max_epochs=(1000 if (max_epochs is None and max_steps is None) else max_epochs),
+        )
+        training_epoch_loop = TrainingEpochLoop(min_steps, max_steps)
+        training_batch_loop = TrainingBatchLoop()
+        validation_epoch_loop = EvaluationLoop()
+        training_epoch_loop.connect(trainer=self, batch_loop=training_batch_loop, val_loop=validation_epoch_loop)
+        training_batch_loop.connect(trainer=self)
+        validation_epoch_loop.connect(trainer=self)
+        self.fit_loop.connect(trainer=self, epoch_loop=training_epoch_loop)
+
+        # .validate() loop
         self.validation_loop = EvaluationLoop()
+        self.validation_loop.connect(trainer=self)
+
+        # .test() loop
         self.test_loop = EvaluationLoop()
+        self.test_loop.connect(trainer=self)
+
+        # .predict() loop
         self.predict_loop = PredictionLoop()
-        self.fit_loop.connect(self)
-        self.validation_loop.connect(self)
-        self.test_loop.connect(self)
-        self.predict_loop.connect(self)
+        self.predict_loop.connect(trainer=self)
 
         # training state
         if weights_summary is not None and weights_summary not in ModelSummary.MODES: