Move results to their files

Author: carmocca

pytorch_lightning/loops/epoch/training_epoch_loop.py

@@ -17,7 +17,7 @@
 
 from pytorch_lightning import loops  # import as loops to avoid circular imports
 from pytorch_lightning.loops.batch import TrainingBatchLoop
-from pytorch_lightning.loops.closure import ClosureResult
+from pytorch_lightning.loops.optimization.closure import OutputResult
 from pytorch_lightning.loops.utilities import _prepare_dataloader_iter
 from pytorch_lightning.trainer.connectors.logger_connector.result import ResultCollection
 from pytorch_lightning.trainer.progress import Progress, SchedulerProgress
@@ -284,18 +284,18 @@ def _track_epoch_end_reduce_metrics(
 
     @staticmethod
     def _prepare_outputs(
-        outputs: List[List[List[ClosureResult]]], batch_mode: bool
+        outputs: List[List[List[OutputResult]]], batch_mode: bool
     ) -> Union[List[List[List[Dict]]], List[List[Dict]], List[Dict], Dict]:
         """Extract required information from batch or epoch end results.
 
         Args:
-            outputs: A 3-dimensional list of ``ClosureResult`` objects with dimensions:
+            outputs: A 3-dimensional list of ``OutputResult`` objects with dimensions:
                 ``[optimizer outs][batch outs][tbptt steps]``.
 
             batch_mode: If True, ignore the batch output dimension.
 
         Returns:
-            The cleaned outputs with ``ClosureResult`` objects converted to dictionaries.
+            The cleaned outputs with ``OutputResult`` objects converted to dictionaries.
             All list dimensions of size one will be collapsed.
         """
         processed_outputs = []
@@ -312,7 +312,7 @@ def _prepare_outputs(
             for batch_outputs in opt_outputs:
                 processed_tbptt_outputs = []
 
-                if isinstance(batch_outputs, ClosureResult):
+                if isinstance(batch_outputs, OutputResult):
                     batch_outputs = [batch_outputs]
 
                 for tbptt_output in batch_outputs:

pytorch_lightning/loops/optimization/closure.py

@@ -12,83 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from abc import ABC, abstractmethod
-from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, Optional
+from dataclasses import dataclass
+from typing import Any, Optional
 
-from torch import Tensor
-
-from pytorch_lightning.profiler import BaseProfiler, PassThroughProfiler
-from pytorch_lightning.utilities.apply_func import apply_to_collection
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities.memory import recursive_detach
-from pytorch_lightning.utilities.types import STEP_OUTPUT
-from pytorch_lightning.utilities.warnings import rank_zero_deprecation, WarningCache
 
 
 @dataclass
-class ClosureResult:
-    """A container to hold the result of a :class:`AbstractClosure` call.
-
-    It is created from the output of :meth:`~pytorch_lightning.core.lightning.LightningModule.training_step`.
-
-    Attributes:
-        closure_loss: The loss with a graph attached.
-        loss: A detached copy of the closure loss.
-        extra: Any keys other than the loss returned.
-    """
-
-    closure_loss: Optional[Tensor]
-    loss: Optional[Tensor] = field(init=False, default=None)
-    extra: Dict[str, Tensor] = field(default_factory=dict)
-
-    def __post_init__(self) -> None:
-        # TODO: remove with the deprecation removal in v1.6
-        ClosureResult._check_extra_detach_deprecation(self.extra)
-        self.extra = recursive_detach(self.extra)
-
-        self._clone_loss()
-
-    def _clone_loss(self) -> None:
-        if self.closure_loss is not None:
-            # the loss will get scaled for amp. avoid any modifications to it
-            self.loss = self.closure_loss.detach().clone()
-
-    @classmethod
-    def from_training_step_output(
-        cls, training_step_output: Optional[STEP_OUTPUT], normalize: int = 1
-    ) -> "ClosureResult":
-        closure_loss, extra = None, {}
-
-        if isinstance(training_step_output, dict):
-            # this should not modify the `training_step_output`, as the user could be using it after `training_step_end`
-            closure_loss = training_step_output.get("loss")
-            extra = {k: v for k, v in training_step_output.items() if k not in ("loss", "hiddens")}
-        elif isinstance(training_step_output, Tensor):
-            closure_loss = training_step_output
-
-        if closure_loss is not None:
-            # accumulate the loss. If ``accumulate_grad_batches == 1``, no effect
-            closure_loss /= normalize
-
-        return cls(closure_loss, extra=extra)
-
-    @staticmethod
-    def _check_extra_detach_deprecation(extra: Dict[str, Any]) -> None:
-        def check_fn(v: Tensor) -> Tensor:
-            if v.grad_fn is not None:
-                rank_zero_deprecation(
-                    f"One of the returned values {set(extra.keys())} has a `grad_fn`. We will detach it automatically"
-                    " but this behaviour will change in v1.6. Please detach it manually:"
-                    " `return {'loss': ..., 'something': something.detach()}`"
-                )
-            return v
-
-        apply_to_collection(extra, Tensor, check_fn)
-
-    def drop_closure_loss(self) -> "ClosureResult":
-        """Return itself without the closure loss which could have a `grad_fn`"""
-        self.closure_loss = None
-        return self
+class OutputResult:
+    ...
 
 
 class AbstractClosure(ABC):
@@ -99,14 +31,14 @@ class AbstractClosure(ABC):
     object which later can call it like a function but without requiring to pass in any arguments.
 
     This class provides a simple abstraction making the instance of this class callable like a function while capturing
-    the :class:`ClosureResult` and caching it.
+    the :class:`OutputResult` and caching it.
     """
 
     def __init__(self) -> None:
         super().__init__()
-        self._result: Optional[ClosureResult] = None
+        self._result: Optional[OutputResult] = None
 
-    def consume_result(self) -> ClosureResult:
+    def consume_result(self) -> OutputResult:
         """The cached result from the last time the closure was called.
 
         Once accessed, the internal reference gets reset and the consumer will have to hold on to the reference as long
@@ -122,69 +54,10 @@ def consume_result(self) -> ClosureResult:
         return result
 
     @abstractmethod
-    def closure(self, *args: Any, **kwargs: Any) -> ClosureResult:
+    def closure(self, *args: Any, **kwargs: Any) -> OutputResult:
         """Implements the behavior of the closure once it is getting called."""
         pass
 
-    def __call__(self, *args: Any, **kwargs: Any) -> Optional[Tensor]:
+    def __call__(self, *args: Any, **kwargs: Any) -> "AbstractClosure":
         self._result = self.closure(*args, **kwargs)
-        return self._result.loss
-
-
-class Closure(AbstractClosure):
-    """An implementation of a :class:`AbstractClosure` for optimization in Lightning that combines three elementary
-    closures into one: ``training_step``, ``backward`` and ``zero_grad``.
-
-    The Closure gets created by the training loop(s) and is then passed to the
-    :meth:`torch.optim.Optimizer.step` method. An optimizer is responsible for calling the closure and optionally
-    do something with the output.
-
-    Args:
-        step_fn: This is typically the :meth:`pytorch_lightning.core.lightning.LightningModule.training_step
-            wrapped with processing for its outputs
-        backward_fn: A function that takes a loss value as input, performs back-propagation and returns the loss value.
-            Can be set to ``None`` to skip the backward operation.
-        zero_grad_fn: A function that zeroes the gradients. Can be set to ``None`` to skip zero_grad, for example
-            when accumulating gradients.
-        profiler: A profiler for profiling the actions of the passed in closure functions.
-
-    Example:
-
-        closure = Closure()
-        optimizer = torch.optim.Adam(...)
-        optimizer.step(closure)
-    """
-
-    warning_cache = WarningCache()
-
-    def __init__(
-        self,
-        step_fn: Callable[[], ClosureResult],
-        backward_fn: Optional[Callable[[Tensor], Tensor]] = None,
-        zero_grad_fn: Optional[Callable[[], None]] = None,
-        profiler: Optional[BaseProfiler] = None,
-    ):
-        super().__init__()
-        self._step_fn = step_fn
-        self._backward_fn = backward_fn
-        self._zero_grad_fn = zero_grad_fn
-        self._profiler = PassThroughProfiler() if profiler is None else profiler
-
-    def closure(self, *args: Any, **kwargs: Any) -> ClosureResult:
-        with self._profiler.profile("training_step_and_backward"):
-            step_output = self._step_fn()
-
-            if step_output.closure_loss is None:
-                self.warning_cache.warn(
-                    "`training_step` returned `None`. If this was on purpose, ignore this warning..."
-                )
-
-            if self._zero_grad_fn is not None:
-                with self._profiler.profile("zero_grad"):
-                    self._zero_grad_fn()
-
-            if self._backward_fn is not None and step_output.closure_loss is not None:
-                with self._profiler.profile("backward"):
-                    step_output.closure_loss = self._backward_fn(step_output.closure_loss)
-
-        return step_output
+        return self

pytorch_lightning/loops/optimization/manual_loop.py

@@ -11,16 +11,89 @@
 # 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 Any, Optional
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+
+from torch import Tensor
 
 from pytorch_lightning.loops import Loop
-from pytorch_lightning.loops.closure import ClosureResult
+from pytorch_lightning.loops.optimization.closure import OutputResult
 from pytorch_lightning.loops.utilities import (
     _build_training_step_kwargs,
     _check_training_step_output,
     _extract_hiddens,
     check_finite_loss,
 )
+from pytorch_lightning.utilities.apply_func import apply_to_collection
+from pytorch_lightning.utilities.memory import recursive_detach
+from pytorch_lightning.utilities.types import STEP_OUTPUT
+from pytorch_lightning.utilities.warnings import rank_zero_deprecation
+
+
+@dataclass
+class ManualResult(OutputResult):
+    """A container to hold the result returned by the ``ManualLoop``.
+
+    It is created from the output of :meth:`~pytorch_lightning.core.lightning.LightningModule.training_step`.
+
+    Attributes:
+        closure_loss: The loss with a graph attached.
+        loss: A detached copy of the closure loss.
+        extra: Any keys other than the loss returned.
+    """
+
+    closure_loss: Optional[Tensor]
+    loss: Optional[Tensor] = field(init=False, default=None)
+    extra: Dict[str, Tensor] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        # TODO: remove with the deprecation removal in v1.6
+        self._check_extra_detach_deprecation(self.extra)
+        self.extra = recursive_detach(self.extra)
+
+        self._clone_loss()
+
+    def _clone_loss(self) -> None:
+        if self.closure_loss is not None:
+            # the loss will get scaled for amp. avoid any modifications to it
+            self.loss = self.closure_loss.detach().clone()
+
+    @classmethod
+    def from_training_step_output(
+        cls, training_step_output: Optional[STEP_OUTPUT], normalize: int = 1
+    ) -> "ManualResult":
+        closure_loss, extra = None, {}
+
+        if isinstance(training_step_output, dict):
+            # this should not modify the `training_step_output`, as the user could be using it after `training_step_end`
+            closure_loss = training_step_output.get("loss")
+            extra = {k: v for k, v in training_step_output.items() if k not in ("loss", "hiddens")}
+        elif isinstance(training_step_output, Tensor):
+            closure_loss = training_step_output
+
+        if closure_loss is not None:
+            # accumulate the loss. If ``accumulate_grad_batches == 1``, no effect
+            closure_loss /= normalize
+
+        return cls(closure_loss, extra=extra)
+
+    @staticmethod
+    def _check_extra_detach_deprecation(extra: Dict[str, Any]) -> None:
+        def check_fn(v: Tensor) -> Tensor:
+            if v.grad_fn is not None:
+                rank_zero_deprecation(
+                    f"One of the returned values {set(extra.keys())} has a `grad_fn`. We will detach it automatically"
+                    " but this behaviour will change in v1.6. Please detach it manually:"
+                    " `return {'loss': ..., 'something': something.detach()}`"
+                )
+            return v
+
+        apply_to_collection(extra, Tensor, check_fn)
+
+    def drop_closure_loss(self) -> "ManualResult":
+        """Return itself without the closure loss which could have a `grad_fn`"""
+        self.closure_loss = None
+        return self
 
 
 class ManualOptimization(Loop):
@@ -36,7 +109,7 @@ def __init__(self) -> None:
         super().__init__()
         self._done: bool = False
         self._hiddens: Optional[Any] = None
-        self._output: Optional[ClosureResult] = None
+        self._output: Optional[ManualResult] = None
 
     @property
     def done(self) -> bool:
@@ -75,8 +148,7 @@ def advance(self, batch: Any, batch_idx: int) -> None:  # type: ignore[override]
 
             self._hiddens = _extract_hiddens(training_step_output, lightning_module.truncated_bptt_steps)
 
-            # TODO: do not use `ClosureResult`
-            result = ClosureResult.from_training_step_output(training_step_output, self.trainer.accumulate_grad_batches)
+            result = ManualResult.from_training_step_output(training_step_output, self.trainer.accumulate_grad_batches)
 
             if self.trainer.terminate_on_nan:
                 check_finite_loss(result.closure_loss)
@@ -90,7 +162,7 @@ def advance(self, batch: Any, batch_idx: int) -> None:  # type: ignore[override]
         self._done = True
         self._output = result
 
-    def on_run_end(self) -> Optional[ClosureResult]:
+    def on_run_end(self) -> Optional[ManualResult]:
         """Returns the result of this loop, i.e., the post-processed outputs from the training step."""
         output, self._output = self._output, None  # free memory
         # #9052 added support for raising `StopIteration` in the `training_step`. If that happens, then `advance`