diff --git a/pyproject.toml b/pyproject.toml index 11f924d51b..707584f2e4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,10 +62,8 @@ ignore_errors = "True" [[tool.mypy.overrides]] module = [ "pytorch_lightning.callbacks.pruning", - "pytorch_lightning.loops.closure", - "pytorch_lightning.loops.batch.manual", - "pytorch_lightning.loops.optimizer", - "pytorch_lightning.trainer.evaluation_loop", + "pytorch_lightning.loops.optimization.*", + "pytorch_lightning.loops.evaluation_loop", "pytorch_lightning.trainer.connectors.logger_connector.*", "pytorch_lightning.trainer.progress", "pytorch_lightning.tuner.auto_gpu_select", diff --git a/pytorch_lightning/loops/__init__.py b/pytorch_lightning/loops/__init__.py index dd0878b3ea..002727274d 100644 --- a/pytorch_lightning/loops/__init__.py +++ b/pytorch_lightning/loops/__init__.py @@ -18,4 +18,4 @@ from pytorch_lightning.loops.batch import TrainingBatchLoop # noqa: F401 from pytorch_lightning.loops.dataloader import DataLoaderLoop, EvaluationLoop, PredictionLoop # noqa: F401 from pytorch_lightning.loops.epoch import EvaluationEpochLoop, PredictionEpochLoop, TrainingEpochLoop # noqa: F401 from pytorch_lightning.loops.fit_loop import FitLoop # noqa: F401 -from pytorch_lightning.loops.optimizer.optimizer_loop import OptimizerLoop # noqa: F401 +from pytorch_lightning.loops.optimization.optimizer_loop import OptimizerLoop # noqa: F401 diff --git a/pytorch_lightning/loops/batch/__init__.py b/pytorch_lightning/loops/batch/__init__.py index 6a5b927ab6..e150fb5881 100644 --- a/pytorch_lightning/loops/batch/__init__.py +++ b/pytorch_lightning/loops/batch/__init__.py @@ -12,5 +12,5 @@ # See the License for the specific language governing permissions and # limitations under the License. -from pytorch_lightning.loops.batch.manual import ManualOptimization # noqa: F401 from pytorch_lightning.loops.batch.training_batch_loop import TrainingBatchLoop # noqa: F401 +from pytorch_lightning.loops.optimization.manual_loop import ManualOptimization # noqa: F401 diff --git a/pytorch_lightning/loops/batch/training_batch_loop.py b/pytorch_lightning/loops/batch/training_batch_loop.py index 4a544b2b9b..8516e0464b 100644 --- a/pytorch_lightning/loops/batch/training_batch_loop.py +++ b/pytorch_lightning/loops/batch/training_batch_loop.py @@ -19,8 +19,8 @@ from torch import Tensor from torch.optim import Optimizer from pytorch_lightning.loops.base import Loop -from pytorch_lightning.loops.batch.manual import ManualOptimization -from pytorch_lightning.loops.optimizer.optimizer_loop import OptimizerLoop +from pytorch_lightning.loops.optimization.manual_loop import ManualOptimization +from pytorch_lightning.loops.optimization.optimizer_loop import OptimizerLoop from pytorch_lightning.trainer.supporters import TensorRunningAccum from pytorch_lightning.utilities import AttributeDict from pytorch_lightning.utilities.types import STEP_OUTPUT diff --git a/pytorch_lightning/loops/epoch/training_epoch_loop.py b/pytorch_lightning/loops/epoch/training_epoch_loop.py index 83fa900931..e0401a2d7b 100644 --- a/pytorch_lightning/loops/epoch/training_epoch_loop.py +++ b/pytorch_lightning/loops/epoch/training_epoch_loop.py @@ -17,7 +17,7 @@ import torch 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 ClosureResult 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 diff --git a/pytorch_lightning/loops/optimizer/__init__.py b/pytorch_lightning/loops/optimization/__init__.py similarity index 86% rename from pytorch_lightning/loops/optimizer/__init__.py rename to pytorch_lightning/loops/optimization/__init__.py index 2f74442741..17e96c49d3 100644 --- a/pytorch_lightning/loops/optimizer/__init__.py +++ b/pytorch_lightning/loops/optimization/__init__.py @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -from pytorch_lightning.loops.optimizer.optimizer_loop import OptimizerLoop # noqa: F401 +from pytorch_lightning.loops.optimization.optimizer_loop import OptimizerLoop # noqa: F401 diff --git a/pytorch_lightning/loops/closure.py b/pytorch_lightning/loops/optimization/closure.py similarity index 100% rename from pytorch_lightning/loops/closure.py rename to pytorch_lightning/loops/optimization/closure.py diff --git a/pytorch_lightning/loops/batch/manual.py b/pytorch_lightning/loops/optimization/manual_loop.py similarity index 98% rename from pytorch_lightning/loops/batch/manual.py rename to pytorch_lightning/loops/optimization/manual_loop.py index f1d933135c..4de3c6af99 100644 --- a/pytorch_lightning/loops/batch/manual.py +++ b/pytorch_lightning/loops/optimization/manual_loop.py @@ -14,7 +14,7 @@ from typing import Any, Optional from pytorch_lightning.loops import Loop -from pytorch_lightning.loops.closure import ClosureResult +from pytorch_lightning.loops.optimization.closure import ClosureResult from pytorch_lightning.loops.utilities import ( _build_training_step_kwargs, _check_training_step_output, diff --git a/pytorch_lightning/loops/optimizer/optimizer_loop.py b/pytorch_lightning/loops/optimization/optimizer_loop.py similarity index 97% rename from pytorch_lightning/loops/optimizer/optimizer_loop.py rename to pytorch_lightning/loops/optimization/optimizer_loop.py index 17f02673ef..aeecedb915 100644 --- a/pytorch_lightning/loops/optimizer/optimizer_loop.py +++ b/pytorch_lightning/loops/optimization/optimizer_loop.py @@ -20,7 +20,7 @@ from torch.optim import Optimizer from pytorch_lightning.core.optimizer import LightningOptimizer from pytorch_lightning.loops import Loop -from pytorch_lightning.loops.closure import Closure, ClosureResult +from pytorch_lightning.loops.optimization.closure import Closure, ClosureResult from pytorch_lightning.loops.utilities import ( _block_parallel_sync_behavior, _build_training_step_kwargs, @@ -43,7 +43,7 @@ class OptimizerLoop(Loop): This loop implements what is known in Lightning as Automatic Optimization. """ - def __init__(self): + def __init__(self) -> None: super().__init__() # TODO: use default dict here to simplify logic in loop self.outputs: _OUTPUTS_TYPE = [] @@ -71,7 +71,7 @@ class OptimizerLoop(Loop): self._batch_idx = batch_idx self._optimizers = optimizers - def advance(self, batch: Any, *args, **kwargs) -> None: # type: ignore[override] + def advance(self, batch: Any, *args: Any, **kwargs: Any) -> None: # type: ignore[override] result = self._run_optimization( batch, self._batch_idx, @@ -183,7 +183,7 @@ class OptimizerLoop(Loop): if not is_first_batch_to_accumulate: return None - def zero_grad_fn(): + def zero_grad_fn() -> None: self._on_before_zero_grad(optimizer) self._optimizer_zero_grad(batch_idx, optimizer, opt_idx) @@ -198,7 +198,7 @@ class OptimizerLoop(Loop): if self._skip_backward: return None - def backward_fn(loss: Tensor): + def backward_fn(loss: Tensor) -> Tensor: self.backward(loss, optimizer, opt_idx) # check if model weights are nan @@ -332,6 +332,7 @@ class OptimizerLoop(Loop): if self.trainer.move_metrics_to_cpu: # hiddens and the training step output are not moved as they are not considered "metrics" + assert self.trainer._results is not None self.trainer._results.cpu() return result diff --git a/tests/core/test_lightning_optimizer.py b/tests/core/test_lightning_optimizer.py index 9d3b4d4dc8..f6809273ff 100644 --- a/tests/core/test_lightning_optimizer.py +++ b/tests/core/test_lightning_optimizer.py @@ -20,7 +20,7 @@ from torch.optim import Adam, Optimizer, SGD from pytorch_lightning import Trainer from pytorch_lightning.core.optimizer import LightningOptimizer -from pytorch_lightning.loops.closure import Closure +from pytorch_lightning.loops.optimization.closure import Closure from tests.helpers.boring_model import BoringModel diff --git a/tests/trainer/loops/__init__.py b/tests/loops/optimization/__init__.py similarity index 100% rename from tests/trainer/loops/__init__.py rename to tests/loops/optimization/__init__.py diff --git a/tests/loops/test_closure.py b/tests/loops/optimization/test_closure.py similarity index 97% rename from tests/loops/test_closure.py rename to tests/loops/optimization/test_closure.py index ff9bd53118..db79d5b4dd 100644 --- a/tests/loops/test_closure.py +++ b/tests/loops/optimization/test_closure.py @@ -15,7 +15,7 @@ import pytest import torch from pytorch_lightning import Trainer -from pytorch_lightning.loops.closure import ClosureResult +from pytorch_lightning.loops.optimization.closure import ClosureResult from pytorch_lightning.utilities.exceptions import MisconfigurationException from tests.helpers import BoringModel diff --git a/tests/trainer/loops/test_all.py b/tests/loops/test_all.py similarity index 100% rename from tests/trainer/loops/test_all.py rename to tests/loops/test_all.py diff --git a/tests/trainer/loops/test_evaluation_loop.py b/tests/loops/test_evaluation_loop.py similarity index 100% rename from tests/trainer/loops/test_evaluation_loop.py rename to tests/loops/test_evaluation_loop.py diff --git a/tests/trainer/loops/test_evaluation_loop_flow.py b/tests/loops/test_evaluation_loop_flow.py similarity index 100% rename from tests/trainer/loops/test_evaluation_loop_flow.py rename to tests/loops/test_evaluation_loop_flow.py diff --git a/tests/trainer/loops/test_flow_warnings.py b/tests/loops/test_flow_warnings.py similarity index 100% rename from tests/trainer/loops/test_flow_warnings.py rename to tests/loops/test_flow_warnings.py diff --git a/tests/trainer/loops/test_training_loop.py b/tests/loops/test_training_loop.py similarity index 100% rename from tests/trainer/loops/test_training_loop.py rename to tests/loops/test_training_loop.py diff --git a/tests/trainer/loops/test_training_loop_flow_dict.py b/tests/loops/test_training_loop_flow_dict.py similarity index 100% rename from tests/trainer/loops/test_training_loop_flow_dict.py rename to tests/loops/test_training_loop_flow_dict.py diff --git a/tests/trainer/loops/test_training_loop_flow_scalar.py b/tests/loops/test_training_loop_flow_scalar.py similarity index 99% rename from tests/trainer/loops/test_training_loop_flow_scalar.py rename to tests/loops/test_training_loop_flow_scalar.py index d5fe6ff6c1..1fb79a267f 100644 --- a/tests/trainer/loops/test_training_loop_flow_scalar.py +++ b/tests/loops/test_training_loop_flow_scalar.py @@ -18,7 +18,7 @@ from torch.utils.data._utils.collate import default_collate from pytorch_lightning import Trainer from pytorch_lightning.core.lightning import LightningModule -from pytorch_lightning.loops.closure import Closure +from pytorch_lightning.loops.optimization.closure import Closure from pytorch_lightning.trainer.states import RunningStage from tests.helpers.boring_model import BoringModel, RandomDataset from tests.helpers.deterministic_model import DeterministicModel