From 2f01c03b38fc16618aa9839d39e0ae5a142c0559 Mon Sep 17 00:00:00 2001 From: schwobr Date: Sat, 7 Dec 2019 14:52:06 +0100 Subject: [PATCH] Additional hooks (#598) * Renamed `on_sanity_check_start` to `on_train_start` and added `on_train_end` to `ModelHooks` * changed tests to use `on_train_start` instead of `on_sanity_check_start` --- pytorch_lightning/core/hooks.py | 16 ++++++++++++++++ pytorch_lightning/trainer/trainer.py | 1 + pytorch_lightning/trainer/training_loop.py | 2 ++ tests/test_restore_models.py | 4 ++-- 4 files changed, 21 insertions(+), 2 deletions(-) diff --git a/pytorch_lightning/core/hooks.py b/pytorch_lightning/core/hooks.py index 496cf153..5d9d9c6f 100644 --- a/pytorch_lightning/core/hooks.py +++ b/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. diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index e21c493d..ccb0c71b 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -489,6 +489,7 @@ class Trainer(TrainerIOMixin, # 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', diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py index d15e6e31..541a99f1 100644 --- a/pytorch_lightning/trainer/training_loop.py +++ b/pytorch_lightning/trainer/training_loop.py @@ -331,6 +331,8 @@ class TrainerTrainLoopMixin(ABC): self.main_progress_bar.close() + model.on_train_end() + if self.logger is not None: self.logger.finalize("success") diff --git a/tests/test_restore_models.py b/tests/test_restore_models.py index 95181b38..b04bb07d 100644 --- a/tests/test_restore_models.py +++ b/tests/test_restore_models.py @@ -247,7 +247,7 @@ def test_dp_resume(tmpdir): # 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 test_cpu_restore_training(tmpdir): 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