Loop Refactor 5/N - Prediction Loop

CHANGELOG.md

@@ -78,7 +78,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
     * Moved attributes `hiddens` and `split_idx` to TrainLoop ([#7507](https://github.com/PyTorchLightning/pytorch-lightning/pull/7507))
     * Refactored the logic around manual and automatic optimization inside the optimizer loop ([#7526](https://github.com/PyTorchLightning/pytorch-lightning/pull/7526))
     * Simplified "should run validation" logic ([#7682](https://github.com/PyTorchLightning/pytorch-lightning/pull/7682))
-    * Refactored "should run validation" logic when the trainer is signaled to stop ([#7701](https://github.com/PyTorchLightning/pytorch-lightning/pull/7701))
+    * Simplified logic for updating the learning rate for schedulers ([#7682](https://github.com/PyTorchLightning/pytorch-lightning/pull/7682))
+    * Removed the `on_epoch` guard from the "should stop" validation check ([#7701](https://github.com/PyTorchLightning/pytorch-lightning/pull/7701))
+    ...
+    * Refactored internal loop interface; added new classes `EpochLoop`, `TrainingLoop`, `BatchLoop` ([#7700](https://github.com/PyTorchLightning/pytorch-lightning/pull/7700))
 
 
 - Refactored logging

pytorch_lightning/callbacks/prediction_writer.py

@@ -109,7 +109,7 @@ def on_predict_batch_end(
         if not self.interval.on_batch:
             return
         is_distributed = trainer.accelerator_connector.is_distributed
-        batch_indices = trainer.predict_loop.batch_indices if is_distributed else None
+        batch_indices = trainer.predict_loop.prediction_loop.current_batch_indices if is_distributed else None
         self.write_on_batch_end(trainer, pl_module, outputs, batch_indices, batch, batch_idx, dataloader_idx)
 
     def on_predict_epoch_end(

pytorch_lightning/callbacks/stochastic_weight_avg.py

@@ -220,12 +220,12 @@ def on_train_epoch_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningMo
             # performing only one pass over the train data-loader to compute activation statistics
             # Therefore, we will virtually increase `num_training_batches` by 1 and skip backward.
             trainer.num_training_batches += 1
-            trainer.train_loop._skip_backward = True
+            trainer.train_loop.skip_backward = True
             self._accumulate_grad_batches = trainer.accumulate_grad_batches
             trainer.accumulate_grad_batches = len(trainer.train_dataloader)
 
     def on_train_epoch_end(self, trainer: 'pl.Trainer', *args):
-        trainer.train_loop._skip_backward = False
+        trainer.train_loop.skip_backward = False
 
     def on_train_end(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule'):
         if self._model_contains_batch_norm and trainer.current_epoch == self.swa_end + 1:

pytorch_lightning/core/lightning.py

@@ -1322,7 +1322,7 @@ def training_step(...):
 
         # backward
         self._running_manual_backward = True
-        self.trainer.train_loop.backward(loss, optimizer=None, opt_idx=None, *args, **kwargs)
+        self.trainer.batch_loop.backward(loss, optimizer=None, opt_idx=None, *args, **kwargs)
         self._running_manual_backward = False
 
     def backward(self, loss: Tensor, optimizer: Optimizer, optimizer_idx: int, *args, **kwargs) -> None:

pytorch_lightning/loops/base.py

@@ -0,0 +1,59 @@
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+from weakref import proxy
+
+import pytorch_lightning as pl
+
+
+class Loop(ABC):
+
+    def __init__(self):
+        self.iteration_count: int = 0
+        self.trainer: Optional['pl.Trainer'] = None
+
+    def connect(self, trainer, *args, **kwargs) -> None:
+        """Connects Loop with all the necessary things like connectors and accelerators"""
+        self.trainer = proxy(trainer)
+
+    @property
+    @abstractmethod
+    def done(self) -> bool:
+        """Property indicating when loop is finished"""
+
+    @abstractmethod
+    def reset(self) -> None:
+        pass
+
+    def run(self, *args: Any, **kwargs: Any) -> Any:
+        self.reset()
+        self.on_run_start(*args, **kwargs)
+
+        while not self.done:
+            try:
+                self.on_advance_start(*args, **kwargs)
+                self.advance(*args, **kwargs)
+                self.on_advance_end()
+                self.iteration_count = self.increment_iteration(self.iteration_count)
+            except StopIteration:
+                break
+
+        return self.on_run_end()
+
+    def on_run_start(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    def on_advance_start(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    @abstractmethod
+    def advance(self, *args: Any, **kwargs: Any) -> None:
+        """What to do within a single step"""
+
+    def on_advance_end(self) -> None:
+        pass
+
+    def on_run_end(self) -> Any:
+        pass
+
+    def increment_iteration(self, iteration: int) -> int:
+        return iteration + 1