diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index cbec5228..07dedd53 100644 --- a/pytorch_lightning/models/trainer.py +++ b/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): diff --git a/pytorch_lightning/root_module/hooks.py b/pytorch_lightning/root_module/hooks.py index 849826a8..06ec614e 100644 --- a/pytorch_lightning/root_module/hooks.py +++ b/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 diff --git a/tests/test_models.py b/tests/test_models.py index 6ee97ffe..14f1e4ac 100644 --- a/tests/test_models.py +++ b/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