Typing callbacks

docs/source/common/optimizers.rst

@@ -35,7 +35,7 @@ To manually optimize, do the following:
   * ``optimizer.step()`` to update your model parameters
 
 Here is a minimal example of manual optimization.
- 
+
 .. testcode:: python
 
     from pytorch_lightning import LightningModule

pl_examples/basic_examples/dali_image_classifier.py

@@ -194,26 +194,15 @@ def __init__(
 
     def train_dataloader(self):
         return DALIClassificationLoader(
-            self.pipe_train,
-            size=len(self.mnist_train),
-            auto_reset=True,
-            fill_last_batch=True
+            self.pipe_train, size=len(self.mnist_train), auto_reset=True, fill_last_batch=True
         )
 
     def val_dataloader(self):
-        return DALIClassificationLoader(
-            self.pipe_val,
-            size=len(self.mnist_val),
-            auto_reset=True,
-            fill_last_batch=False
-        )
+        return DALIClassificationLoader(self.pipe_val, size=len(self.mnist_val), auto_reset=True, fill_last_batch=False)
 
     def test_dataloader(self):
         return DALIClassificationLoader(
-            self.pipe_test,
-            size=len(self.mnist_test),
-            auto_reset=True,
-            fill_last_batch=False
+            self.pipe_test, size=len(self.mnist_test), auto_reset=True, fill_last_batch=False
         )
 
 

pl_examples/basic_examples/profiler_example.py

@@ -44,10 +44,7 @@
 
 class ModelToProfile(LightningModule):
 
-    def __init__(
-        self,
-        name: str = "resnet50"
-    ):
+    def __init__(self, name: str = "resnet50"):
         super().__init__()
         self.model = getattr(models, name)(pretrained=True)
         self.criterion = torch.nn.CrossEntropyLoss()

pytorch_lightning/__about__.py

@@ -1,7 +1,7 @@
 import time
 
 _this_year = time.strftime("%Y")
-__version__ = '1.3.0rc2'
+__version__ = "20210420"
 __author__ = 'William Falcon et al.'
 __author_email__ = '[email protected]'
 __license__ = 'Apache-2.0'

pytorch_lightning/accelerators/accelerator.py

@@ -24,14 +24,11 @@
 from pytorch_lightning.plugins.precision import ApexMixedPrecisionPlugin, NativeMixedPrecisionPlugin, PrecisionPlugin
 from pytorch_lightning.plugins.training_type import TrainingTypePlugin
 from pytorch_lightning.trainer.states import TrainerState
-from pytorch_lightning.utilities import _NATIVE_AMP_AVAILABLE, rank_zero_warn
+from pytorch_lightning.utilities import rank_zero_warn
 from pytorch_lightning.utilities.apply_func import move_data_to_device
-from pytorch_lightning.utilities.enums import AMPType, GradClipAlgorithmType, LightningEnum
+from pytorch_lightning.utilities.enums import AMPType, GradClipAlgorithmType
 from pytorch_lightning.utilities.types import EPOCH_OUTPUT, STEP_OUTPUT
 
-if _NATIVE_AMP_AVAILABLE:
-    from torch.cuda.amp import GradScaler
-
 
 class Accelerator:
     """
@@ -374,7 +371,7 @@ def to_device(self, batch: Any) -> Any:
         return self.batch_to_device(batch, self.root_device)
 
     @property
-    def amp_backend(self) -> Optional[LightningEnum]:
+    def amp_backend(self) -> Optional['AMPType']:
         if isinstance(self.precision_plugin, ApexMixedPrecisionPlugin):
             return AMPType.APEX
         elif isinstance(self.precision_plugin, NativeMixedPrecisionPlugin):
@@ -386,7 +383,7 @@ def precision(self) -> Union[str, int]:
         return self.precision_plugin.precision
 
     @property
-    def scaler(self) -> Optional['GradScaler']:
+    def scaler(self) -> Optional['torch.cuda.amp.GradScaler']:
         return getattr(self.precision_plugin, 'scaler', None)
 
     @property

pytorch_lightning/accelerators/cpu.py

@@ -11,6 +11,8 @@
 # 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 TYPE_CHECKING
+
 import pytorch_lightning as pl
 from pytorch_lightning.accelerators.accelerator import Accelerator
 from pytorch_lightning.plugins.precision import MixedPrecisionPlugin

pytorch_lightning/accelerators/gpu.py

@@ -13,17 +13,18 @@
 # limitations under the License.
 import logging
 import os
-from typing import Any
+from typing import Any, TYPE_CHECKING
 
 import torch
 
-import pytorch_lightning as pl
 from pytorch_lightning.accelerators.accelerator import Accelerator
 from pytorch_lightning.plugins import DataParallelPlugin
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
 _log = logging.getLogger(__name__)
 
+import pytorch_lightning as pl
+
 
 class GPUAccelerator(Accelerator):
     """ Accelerator for GPU devices. """

pytorch_lightning/accelerators/tpu.py

@@ -11,7 +11,7 @@
 # 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, Callable, Union
+from typing import Any, Callable, TYPE_CHECKING, Union
 
 from torch.optim import Optimizer
 
@@ -51,7 +51,7 @@ def setup(self, trainer: 'pl.Trainer', model: 'pl.LightningModule') -> None:
         return super().setup(trainer, model)
 
     def run_optimizer_step(
-        self, optimizer: Optimizer, optimizer_idx: int, lambda_closure: Callable, **kwargs: Any
+        self, optimizer: 'Optimizer', optimizer_idx: int, lambda_closure: Callable, **kwargs: Any
     ) -> None:
         xm.optimizer_step(optimizer, barrier=False, optimizer_args={'closure': lambda_closure, **kwargs})
 

pytorch_lightning/callbacks/base.py

@@ -17,9 +17,11 @@
 """
 
 import abc
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
 
-from pytorch_lightning.core.lightning import LightningModule
+from torch.optim import Optimizer
+
+import pytorch_lightning as pl
 
 
 class Callback(abc.ABC):
@@ -29,158 +31,165 @@ class Callback(abc.ABC):
     Subclass this class and override any of the relevant hooks
     """
 
-    def on_configure_sharded_model(self, trainer, pl_module: LightningModule) -> None:
+    def on_configure_sharded_model(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called before configure sharded model"""
 
-    def on_before_accelerator_backend_setup(self, trainer, pl_module: LightningModule) -> None:
+    def on_before_accelerator_backend_setup(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called before accelerator is being setup"""
         pass
 
-    def setup(self, trainer, pl_module: LightningModule, stage: Optional[str] = None) -> None:
+    def setup(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', stage: Optional[str] = None) -> None:
         """Called when fit, validate, test, predict, or tune begins"""
         pass
 
-    def teardown(self, trainer, pl_module: LightningModule, stage: Optional[str] = None) -> None:
+    def teardown(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', stage: Optional[str] = None) -> None:
         """Called when fit, validate, test, predict, or tune ends"""
         pass
 
-    def on_init_start(self, trainer) -> None:
+    def on_init_start(self, trainer: 'pl.Trainer') -> None:
         """Called when the trainer initialization begins, model has not yet been set."""
         pass
 
-    def on_init_end(self, trainer) -> None:
+    def on_init_end(self, trainer: 'pl.Trainer') -> None:
         """Called when the trainer initialization ends, model has not yet been set."""
         pass
 
-    def on_fit_start(self, trainer, pl_module: LightningModule) -> None:
+    def on_fit_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when fit begins"""
         pass
 
-    def on_fit_end(self, trainer, pl_module: LightningModule) -> None:
+    def on_fit_end(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when fit ends"""
         pass
 
-    def on_sanity_check_start(self, trainer, pl_module: LightningModule) -> None:
+    def on_sanity_check_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when the validation sanity check starts."""
         pass
 
-    def on_sanity_check_end(self, trainer, pl_module: LightningModule) -> None:
+    def on_sanity_check_end(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when the validation sanity check ends."""
         pass
 
     def on_train_batch_start(
-        self, trainer, pl_module: LightningModule, batch: Any, batch_idx: int, dataloader_idx: int
+        self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', batch: Any, batch_idx: int, dataloader_idx: int
     ) -> None:
         """Called when the train batch begins."""
         pass
 
     def on_train_batch_end(
-        self, trainer, pl_module: LightningModule, outputs: Any, batch: Any, batch_idx: int, dataloader_idx: int
+        self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', outputs: Any, batch: Any, batch_idx: int,
+        dataloader_idx: int
     ) -> None:
         """Called when the train batch ends."""
         pass
 
-    def on_train_epoch_start(self, trainer, pl_module: LightningModule) -> None:
+    def on_train_epoch_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when the train epoch begins."""
         pass
 
-    def on_train_epoch_end(self, trainer, pl_module: LightningModule, outputs: List[Any]) -> None:
+    def on_train_epoch_end(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', outputs: List[Any]) -> None:
         """Called when the train epoch ends."""
         pass
 
-    def on_validation_epoch_start(self, trainer, pl_module: LightningModule) -> None:
+    def on_validation_epoch_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when the val epoch begins."""
         pass
 
-    def on_validation_epoch_end(self, trainer, pl_module: LightningModule, outputs: List[Any]) -> None:
+    def on_validation_epoch_end(
+        self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', outputs: List[Any]
+    ) -> None:
         """Called when the val epoch ends."""
         pass
 
-    def on_test_epoch_start(self, trainer, pl_module: LightningModule) -> None:
+    def on_test_epoch_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when the test epoch begins."""
         pass
 
-    def on_test_epoch_end(self, trainer, pl_module: LightningModule, outputs: List[Any]) -> None:
+    def on_test_epoch_end(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', outputs: List[Any]) -> None:
         """Called when the test epoch ends."""
         pass
 
-    def on_epoch_start(self, trainer, pl_module: LightningModule) -> None:
+    def on_epoch_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when either of train/val/test epoch begins."""
         pass
 
-    def on_epoch_end(self, trainer, pl_module: LightningModule) -> None:
+    def on_epoch_end(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when either of train/val/test epoch ends."""
         pass
 
-    def on_batch_start(self, trainer, pl_module: LightningModule) -> None:
+    def on_batch_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when the training batch begins."""
         pass
 
     def on_validation_batch_start(
-        self, trainer, pl_module: LightningModule, batch: Any, batch_idx: int, dataloader_idx: int
+        self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', batch: Any, batch_idx: int, dataloader_idx: int
     ) -> None:
         """Called when the validation batch begins."""
         pass
 
     def on_validation_batch_end(
-        self, trainer, pl_module: LightningModule, outputs: Any, batch: Any, batch_idx: int, dataloader_idx: int
+        self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', outputs: Any, batch: Any, batch_idx: int,
+        dataloader_idx: int
     ) -> None:
         """Called when the validation batch ends."""
         pass
 
     def on_test_batch_start(
-        self, trainer, pl_module: LightningModule, batch: Any, batch_idx: int, dataloader_idx: int
+        self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', batch: Any, batch_idx: int, dataloader_idx: int
     ) -> None:
         """Called when the test batch begins."""
         pass
 
     def on_test_batch_end(
-        self, trainer, pl_module: LightningModule, outputs: Any, batch: Any, batch_idx: int, dataloader_idx: int
+        self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', outputs: Any, batch: Any, batch_idx: int,
+        dataloader_idx: int
     ) -> None:
         """Called when the test batch ends."""
         pass
 
-    def on_batch_end(self, trainer, pl_module: LightningModule) -> None:
+    def on_batch_end(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when the training batch ends."""
         pass
 
-    def on_train_start(self, trainer, pl_module: LightningModule) -> None:
+    def on_train_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when the train begins."""
         pass
 
-    def on_train_end(self, trainer, pl_module: LightningModule) -> None:
+    def on_train_end(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when the train ends."""
         pass
 
-    def on_pretrain_routine_start(self, trainer, pl_module: LightningModule) -> None:
+    def on_pretrain_routine_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when the pretrain routine begins."""
         pass
 
-    def on_pretrain_routine_end(self, trainer, pl_module: LightningModule) -> None:
+    def on_pretrain_routine_end(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when the pretrain routine ends."""
         pass
 
-    def on_validation_start(self, trainer, pl_module: LightningModule) -> None:
+    def on_validation_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when the validation loop begins."""
         pass
 
-    def on_validation_end(self, trainer, pl_module: LightningModule) -> None:
+    def on_validation_end(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when the validation loop ends."""
         pass
 
-    def on_test_start(self, trainer, pl_module: LightningModule) -> None:
+    def on_test_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when the test begins."""
         pass
 
-    def on_test_end(self, trainer, pl_module: LightningModule) -> None:
+    def on_test_end(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when the test ends."""
         pass
 
-    def on_keyboard_interrupt(self, trainer, pl_module: LightningModule) -> None:
+    def on_keyboard_interrupt(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called when the training is interrupted by ``KeyboardInterrupt``."""
         pass
 
-    def on_save_checkpoint(self, trainer, pl_module: LightningModule, checkpoint: Dict[str, Any]) -> dict:
+    def on_save_checkpoint(
+        self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', checkpoint: Dict[str, Any]
+    ) -> dict:
         """
         Called when saving a model checkpoint, use to persist state.
 
@@ -202,10 +211,12 @@ def on_load_checkpoint(self, callback_state: Dict[str, Any]) -> None:
         """
         pass
 
-    def on_after_backward(self, trainer, pl_module: LightningModule) -> None:
+    def on_after_backward(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called after ``loss.backward()`` and before optimizers do anything."""
         pass
 
-    def on_before_zero_grad(self, trainer, pl_module: LightningModule, optimizer) -> None:
+    def on_before_zero_grad(
+        self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', optimizer: 'Optimizer'
+    ) -> None:
         """Called after ``optimizer.step()`` and before ``optimizer.zero_grad()``."""
         pass