Checkpointing interval (#1272)

* formatting

* formatting

* fix interval

* fix train loop

* fix test

* parametrize test

* Apply suggestions from code review

Co-Authored-By: Adrian Wälchli <adrian.waelchli@students.unibe.ch>

* fix calling

* flake8

* add types

Co-authored-by: Adrian Wälchli <adrian.waelchli@students.unibe.ch>
Co-authored-by: William Falcon <waf2107@columbia.edu>

tests/base/__init__.py

@@ -40,7 +40,7 @@ class LightningTestModel(LightTrainDataloader,
 
 
 class LightningTestModelWithoutHyperparametersArg(LightningTestModel):
-    """ without hparams argument in constructor """
+    """Without hparams argument in constructor """
 
     def __init__(self):
         import tests.base.utils as tutils
@@ -51,7 +51,7 @@ class LightningTestModelWithoutHyperparametersArg(LightningTestModel):
 
 
 class LightningTestModelWithUnusedHyperparametersArg(LightningTestModelWithoutHyperparametersArg):
-    """ has hparams argument in constructor but is not used """
+    """It has hparams argument in constructor but is not used."""
 
     def __init__(self, hparams):
         super().__init__()

tests/base/mixins.py

@@ -14,11 +14,7 @@ class LightValidationStepMixin:
         return self._dataloader(train=False)
 
     def validation_step(self, batch, batch_idx, *args, **kwargs):
-        """
-        Lightning calls this inside the validation loop
-        :param batch:
-        :return:
-        """
+        """Lightning calls this inside the validation loop."""
         x, y = batch
         x = x.view(x.size(0), -1)
         y_hat = self(x)
@@ -66,8 +62,9 @@ class LightValidationMixin(LightValidationStepMixin):
     def validation_epoch_end(self, outputs):
         """
         Called at the end of validation to aggregate outputs
-        :param outputs: list of individual outputs of each validation step
-        :return:
+
+        Args:
+            outputs: list of individual outputs of each validation step
         """
         # if returned a scalar from validation_step, outputs is a list of tensor scalars
         # we return just the average in this case (if we want)

tests/base/models.py

@@ -42,16 +42,10 @@ class DictHparamsModel(LightningModule):
 
 
 class TestModelBase(LightningModule):
-    """
-    Base LightningModule for testing. Implements only the required
-    interface
-    """
+    """Base LightningModule for testing. Implements only the required interface."""
 
-    def __init__(self, hparams, force_remove_distributed_sampler=False):
-        """
-        Pass in parsed HyperOptArgumentParser to the model
-        :param hparams:
-        """
+    def __init__(self, hparams, force_remove_distributed_sampler: bool = False):
+        """Pass in parsed HyperOptArgumentParser to the model."""
         # init superclass
         super().__init__()
         self.hparams = hparams
@@ -71,10 +65,7 @@ class TestModelBase(LightningModule):
     # MODEL SETUP
     # ---------------------
     def __build_model(self):
-        """
-        Layout model
-        :return:
-        """
+        """Layout model."""
         self.c_d1 = nn.Linear(in_features=self.hparams.in_features,
                               out_features=self.hparams.hidden_dim)
         self.c_d1_bn = nn.BatchNorm1d(self.hparams.hidden_dim)
@@ -87,11 +78,7 @@ class TestModelBase(LightningModule):
     # TRAINING
     # ---------------------
     def forward(self, x):
-        """
-        No special modification required for lightning, define as you normally would
-        :param x:
-        :return:
-        """
+        """No special modification required for lightning, define as you normally would."""
         x = self.c_d1(x)
         x = torch.tanh(x)
         x = self.c_d1_bn(x)
@@ -107,11 +94,7 @@ class TestModelBase(LightningModule):
         return nll
 
     def training_step(self, batch, batch_idx, optimizer_idx=None):
-        """
-        Lightning calls this inside the training loop
-        :param batch:
-        :return:
-        """
+        """Lightning calls this inside the training loop"""
         # forward pass
         x, y = batch
         x = x.view(x.size(0), -1)