added hook on_sanity_check_start

pytorch_lightning/models/trainer.py

@@ -621,9 +621,6 @@ class Trainer(TrainerIO):
         ref_model.trainer = self
         ref_model.experiment = self.experiment
 
-        # run tiny validation to make sure program won't crash during val
-        _ = self.validate(model, self.val_dataloader, max_batches=self.nb_sanity_val_steps)
-
         # save exp to get started
         if self.proc_rank == 0:
             self.experiment.save()
@@ -641,9 +638,14 @@ class Trainer(TrainerIO):
         if self.cluster is not None:  # pragma: no cover
             self.enable_auto_hpc_walltime_manager()
 
+        # run tiny validation to make sure program won't crash during val
+        model.on_sanity_check_start()
+        _ = self.validate(model, self.val_dataloader, max_batches=self.nb_sanity_val_steps)
+
         # ---------------------------
         # CORE TRAINING LOOP
         # ---------------------------
+
         self.__train()
 
     def __train(self):

pytorch_lightning/root_module/hooks.py

@@ -2,6 +2,14 @@ import torch
 
 
 class ModelHooks(torch.nn.Module):
+
+    def on_sanity_check_start(self):
+        """
+        Called before starting validate
+        :return:
+        """
+        pass
+
     def on_batch_start(self, data_batch):
         pass
 

tests/test_models.py

@@ -93,7 +93,7 @@ def test_cpu_restore_training():
         new_pred = trainer.model(x)
         assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
 
-    model.on_epoch_start = assert_pred_same
+    model.on_sanity_check_start = assert_pred_same
 
     # 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