diff --git a/CHANGELOG.md b/CHANGELOG.md index 9a4abf77b2a17..6bf80b1ccf739 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -47,6 +47,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Added `IoU` class interface ([#4704](https://github.com/PyTorchLightning/pytorch-lightning/pull/4704)) +- Added missing val/test hooks in `LightningModule` ([#5467](https://github.com/PyTorchLightning/pytorch-lightning/pull/5467)) + + ### Changed - Changed `stat_scores` metric now calculates stat scores over all classes and gains new parameters, in line with the new `StatScores` metric ([#4839](https://github.com/PyTorchLightning/pytorch-lightning/pull/4839)) diff --git a/pytorch_lightning/core/hooks.py b/pytorch_lightning/core/hooks.py index f24a4ce8beb8a..0b97d328904ac 100644 --- a/pytorch_lightning/core/hooks.py +++ b/pytorch_lightning/core/hooks.py @@ -17,14 +17,15 @@ from typing import Any, Dict, List, Optional, Union import torch -from pytorch_lightning.utilities import move_data_to_device, rank_zero_warn from torch.optim.optimizer import Optimizer from torch.utils.data import DataLoader +from pytorch_lightning.utilities import move_data_to_device, rank_zero_warn + class ModelHooks: """Hooks to be used in LightningModule.""" - def setup(self, stage: str): + def setup(self, stage: str) -> None: """ Called at the beginning of fit and test. This is a good hook when you need to build models dynamically or adjust something about them. @@ -52,7 +53,7 @@ def setup(stage): """ - def teardown(self, stage: str): + def teardown(self, stage: str) -> None: """ Called at the end of fit and test. @@ -60,13 +61,13 @@ def teardown(self, stage: str): stage: either 'fit' or 'test' """ - def on_fit_start(self): + def on_fit_start(self) -> None: """ Called at the very beginning of fit. If on DDP it is called on every process """ - def on_fit_end(self): + def on_fit_end(self) -> None: """ Called at the very end of fit. If on DDP it is called on every process @@ -74,7 +75,7 @@ def on_fit_end(self): def on_train_start(self) -> None: """ - Called at the beginning of training before sanity check. + Called at the beginning of training after sanity check. """ # do something at the start of training @@ -84,6 +85,18 @@ def on_train_end(self) -> None: """ # do something at the end of training + def on_validation_start(self) -> None: + """ + Called at the beginning of validation. + """ + # do something at the start of validation + + def on_validation_end(self) -> None: + """ + Called at the end of validation. + """ + # do something at the end of validation + def on_pretrain_routine_start(self) -> None: """ Called at the beginning of the pretrain routine (between fit and train start). @@ -108,9 +121,7 @@ def on_pretrain_routine_end(self) -> None: """ # do something at the end of the pretrain routine - def on_train_batch_start( - self, batch: Any, batch_idx: int, dataloader_idx: int - ) -> None: + def on_train_batch_start(self, batch: Any, batch_idx: int, dataloader_idx: int) -> None: """ Called in the training loop before anything happens for that batch. @@ -253,6 +264,18 @@ def on_test_epoch_end(self) -> None: """ # do something when the epoch ends + def on_test_start(self) -> None: + """ + Called at the beginning of testing. + """ + # do something at the start of testing + + def on_test_end(self) -> None: + """ + Called at the end of testing. + """ + # do something at the end of testing + def on_before_zero_grad(self, optimizer: Optimizer) -> None: """ Called after optimizer.step() and before optimizer.zero_grad(). diff --git a/tests/models/test_hooks.py b/tests/models/test_hooks.py index 250d50a9b03ba..1f25d46f82944 100644 --- a/tests/models/test_hooks.py +++ b/tests/models/test_hooks.py @@ -12,15 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. import inspect +from unittest.mock import MagicMock import pytest import torch -from unittest.mock import MagicMock from pytorch_lightning import Trainer from pytorch_lightning.accelerators.gpu_accelerator import GPUAccelerator from pytorch_lightning.trainer.states import TrainerState -from tests.base import EvalModelTemplate, BoringModel +from tests.base import BoringModel, EvalModelTemplate @pytest.mark.parametrize('max_steps', [1, 2, 3]) @@ -254,10 +254,6 @@ def on_test_start(self): self.called.append(inspect.currentframe().f_code.co_name) super().on_test_start() - def on_test_end(self): - self.called.append(inspect.currentframe().f_code.co_name) - super().on_test_end() - def on_test_batch_start(self, batch, batch_idx, dataloader_idx): self.called.append(inspect.currentframe().f_code.co_name) super().on_test_batch_start(batch, batch_idx, dataloader_idx) @@ -290,6 +286,14 @@ def on_test_model_train(self): self.called.append(inspect.currentframe().f_code.co_name) super().on_test_model_train() + def on_test_end(self): + self.called.append(inspect.currentframe().f_code.co_name) + super().on_test_end() + + def teardown(self, stage: str): + self.called.append(inspect.currentframe().f_code.co_name) + super().teardown(stage) + model = HookedModel() assert model.called == [] @@ -313,10 +317,12 @@ def on_test_model_train(self): 'on_pretrain_routine_start', 'on_pretrain_routine_end', 'on_validation_model_eval', + 'on_validation_start', 'on_validation_epoch_start', 'on_validation_batch_start', 'on_validation_batch_end', 'on_validation_epoch_end', + 'on_validation_end', 'on_validation_model_train', 'on_train_start', 'on_epoch_start', @@ -330,16 +336,19 @@ def on_test_model_train(self): 'on_before_zero_grad', 'on_train_batch_end', 'on_validation_model_eval', + 'on_validation_start', 'on_validation_epoch_start', 'on_validation_batch_start', 'on_validation_batch_end', 'on_validation_epoch_end', 'on_save_checkpoint', + 'on_validation_end', 'on_validation_model_train', 'on_epoch_end', 'on_train_epoch_end', 'on_train_end', 'on_fit_end', + 'teardown', ] assert model.called == expected @@ -352,12 +361,16 @@ def on_test_model_train(self): 'on_pretrain_routine_start', 'on_pretrain_routine_end', 'on_test_model_eval', + 'on_test_start', 'on_test_epoch_start', 'on_test_batch_start', 'on_test_batch_end', 'on_test_epoch_end', + 'on_test_end', 'on_test_model_train', 'on_fit_end', + 'teardown', # for 'fit' + 'teardown', # for 'test' ] assert model2.called == expected