Add typing to lightning.tuner

pytorch_lightning/tuner/batch_size_scaling.py

@@ -251,5 +251,5 @@ def _adjust_batch_size(
     return new_size, changed
 
 
-def _is_valid_batch_size(current_size, dataloader):
+def _is_valid_batch_size(current_size: int, dataloader) -> bool:
     return not has_len(dataloader) or current_size <= len(dataloader)

pytorch_lightning/tuner/lr_finder.py

@@ -98,7 +98,7 @@ def __init__(self, mode: str, lr_min: float, lr_max: float, num_training: int):
         self.results = {}
         self._total_batch_idx = 0  # for debug purpose
 
-    def _exchange_scheduler(self, configure_optimizers: Callable):
+    def _exchange_scheduler(self, configure_optimizers: Callable) -> Callable:
         """ Decorate configure_optimizers methods such that it returns the users
             originally specified optimizer together with a new scheduler that
             that takes care of the learning rate search.
@@ -280,7 +280,7 @@ def lr_find(
     return lr_finder
 
 
-def __lr_finder_dump_params(trainer, model):
+def __lr_finder_dump_params(trainer: 'pl.Trainer', model: 'pl.LightningModule') -> None:
     # Prevent going into infinite loop
     trainer.__dumped_params = {
         'auto_lr_find': trainer.auto_lr_find,
@@ -292,7 +292,7 @@ def __lr_finder_dump_params(trainer, model):
     }
 
 
-def __lr_finder_restore_params(trainer, model):
+def __lr_finder_restore_params(trainer: 'pl.Trainer', model: 'pl.LightningModule') -> None:
     trainer.auto_lr_find = trainer.__dumped_params['auto_lr_find']
     trainer.logger = trainer.__dumped_params['logger']
     trainer.callbacks = trainer.__dumped_params['callbacks']
@@ -336,7 +336,7 @@ def __init__(
         self.progress_bar_refresh_rate = progress_bar_refresh_rate
         self.progress_bar = None
 
-    def on_batch_start(self, trainer, pl_module):
+    def on_batch_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """ Called before each training batch, logs the lr that will be used """
         if (trainer.batch_idx + 1) % trainer.accumulate_grad_batches != 0:
             return
@@ -346,7 +346,10 @@ def on_batch_start(self, trainer, pl_module):
 
         self.lrs.append(trainer.lr_schedulers[0]['scheduler'].lr[0])
 
-    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
+    def on_train_batch_end(
+        self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', outputs, batch, batch_idx: Optional[int],
+        dataloader_idx: Optional[int]
+    ) -> None:
         """ Called when the training batch ends, logs the calculated loss """
         if (trainer.batch_idx + 1) % trainer.accumulate_grad_batches != 0:
             return
@@ -397,7 +400,7 @@ def __init__(self, optimizer: torch.optim.Optimizer, end_lr: float, num_iter: in
         self.num_iter = num_iter
         super(_LinearLR, self).__init__(optimizer, last_epoch)
 
-    def get_lr(self):
+    def get_lr(self) -> list:
         curr_iter = self.last_epoch + 1
         r = curr_iter / self.num_iter
 
@@ -435,7 +438,7 @@ def __init__(self, optimizer: torch.optim.Optimizer, end_lr: float, num_iter: in
         self.num_iter = num_iter
         super(_ExponentialLR, self).__init__(optimizer, last_epoch)
 
-    def get_lr(self):
+    def get_lr(self) -> list:
         curr_iter = self.last_epoch + 1
         r = curr_iter / self.num_iter