[feat] Add restore to base loop

CHANGELOG.md

@@ -130,6 +130,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `max_depth` parameter in `ModelSummary` ([#8062](https://github.com/PyTorchLightning/pytorch-lightning/pull/8062))
 
 
+- Added `restore` function and `restarting` attribute to base `Loop` ([#8247](https://github.com/PyTorchLightning/pytorch-lightning/pull/8247))
+
+
 ### Changed
 
 

pytorch_lightning/loops/base.py

@@ -46,6 +46,15 @@ class Loop(ABC):
     def __init__(self) -> None:
         self.iteration_count: int = 0
         self.trainer: Optional['pl.Trainer'] = None
+        self._restarting = False
+
+    @property
+    def restarting(self) -> bool:
+        return self._restarting
+
+    @restarting.setter
+    def restarting(self, restarting: bool) -> None:
+        self._restarting = restarting
 
     @property
     @abstractmethod
@@ -87,7 +96,12 @@ def run(self, *args: Any, **kwargs: Any) -> Optional[Any]:
         if self.skip:
             return self.on_skip()
 
-        self.reset()
+        if self.restarting:
+            self.restore()
+            self.restarting = False
+        else:
+            self.reset()
+
         self.on_run_start(*args, **kwargs)
 
         while not self.done:
@@ -103,6 +117,9 @@ def run(self, *args: Any, **kwargs: Any) -> Optional[Any]:
         self.teardown()
         return output
 
+    def restore(self) -> None:
+        """Restore the internal state of the loop the beginning of run if restarting is ``True``."""
+
     @abstractmethod
     def reset(self) -> None:
         """Resets the internal state of the loop at the beginning of each call to :attr:`run`."""

tests/loops/test_loops.py

@@ -0,0 +1,74 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, Iterator
+
+from pytorch_lightning.loops.base import Loop
+
+
+def test_loop_restore():
+
+    class CustomExpection(Exception):
+        pass
+
+    class Simple(Loop):
+
+        def __init__(self, dataset: Iterator):
+            super().__init__()
+            self.dataset = dataset
+
+        def restore(self) -> None:
+            self.iter_dataset = iter(self.dataset)
+            for _ in range(self.iteration_count):
+                next(self.iter_dataset)
+            self.iteration_count += 1
+
+        @property
+        def done(self) -> bool:
+            return self.iteration_count > len(self.dataset)
+
+        def reset(self) -> None:
+            self.iter_dataset = iter(self.dataset)
+            self.outputs = []
+
+        def advance(self) -> None:
+            value = next(self.iter_dataset)
+
+            if self.iteration_count == 5:
+                raise CustomExpection
+
+            self.outputs.append(value)
+
+        def state_dict(self) -> Dict:
+            return {"iteration_count": self.iteration_count, "outputs": self.outputs}
+
+        def load_state_dict(self, state_dict: Dict) -> None:
+            self.iteration_count = state_dict["iteration_count"]
+            self.outputs = state_dict["outputs"]
+
+    data = range(10)
+    loop = Simple(data)
+    try:
+        loop.run()
+        state_dict = {}
+    except CustomExpection:
+        state_dict = loop.state_dict()
+
+    loop = Simple(data)
+    loop.load_state_dict(state_dict)
+    loop.restarting = True
+    loop.run()
+
+    assert not loop.restarting
+    assert loop.outputs == list(range(10))