Additional hooks

pytorch_lightning/core/hooks.py

@@ -28,10 +28,26 @@ class ModelHooks(torch.nn.Module):
     def on_sanity_check_start(self):
         """
         Called before starting evaluate
+        .. warning:: will be deprecated.
         :return:
         """
         pass
 
+    def on_train_start(self):
+        """Called at the beginning of training before sanity check
+        :return:
+        """
+        # do something at the start of training
+        pass
+
+    def on_train_end(self):
+        """
+        Called at the end of training before logger experiment is closed
+        :return:
+        """
+        # do something at the end of training
+        pass
+
     def on_batch_start(self, batch):
         """Called in the training loop before anything happens for that batch.
 

pytorch_lightning/trainer/trainer.py

@@ -489,6 +489,7 @@ def run_pretrain_routine(self, model):
         # run tiny validation (if validation defined)
         # to make sure program won't crash during val
         ref_model.on_sanity_check_start()
+        ref_model.on_train_start()
         if self.get_val_dataloaders() is not None and self.num_sanity_val_steps > 0:
             # init progress bars for validation sanity check
             pbar = tqdm.tqdm(desc='Validation sanity check', total=self.num_sanity_val_steps,

pytorch_lightning/trainer/training_loop.py

@@ -331,6 +331,8 @@ def train(self):
 
         self.main_progress_bar.close()
 
+        model.on_train_end()
+
         if self.logger is not None:
             self.logger.finalize("success")
 

tests/test_restore_models.py

@@ -247,7 +247,7 @@ def assert_good_acc():
 
     # new model
     model = LightningTestModel(hparams)
-    model.on_sanity_check_start = assert_good_acc
+    model.on_train_start = assert_good_acc
 
     # fit new model which should load hpc weights
     new_trainer.fit(model)
@@ -311,7 +311,7 @@ def assert_good_acc():
         for dataloader in trainer.get_val_dataloaders():
             tutils.run_prediction(dataloader, trainer.model)
 
-    model.on_sanity_check_start = assert_good_acc
+    model.on_train_start = assert_good_acc
 
     # by calling fit again, we trigger training, loading weights from the cluster
     # and our hook to predict using current model before any more weight updates