diff --git a/CHANGELOG.md b/CHANGELOG.md index f3190f9d..46c2ec88 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Added type hints in `Trainer.fit()` and `Trainer.test()` to reflect that also a list of dataloaders can be passed in ([#1723](https://github.com/PyTorchLightning/pytorch-lightning/pull/1723)). +- The progress bar metrics now also get updated in `training_epoch_end` ([#1724](https://github.com/PyTorchLightning/pytorch-lightning/pull/1724)). + ### Changed - Reduction when `batch_size < num_gpus` ([#1609](https://github.com/PyTorchLightning/pytorch-lightning/pull/1609)) diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py index 1b906c0f..8638a512 100644 --- a/pytorch_lightning/core/lightning.py +++ b/pytorch_lightning/core/lightning.py @@ -257,6 +257,7 @@ class LightningModule(ABC, GradInformation, ModelIO, ModelHooks): May contain the following optional keys: - log (metrics to be added to the logger; only tensors) + - progress_bar (dict for progress bar display) - any metric used in a callback (e.g. early stopping). Note: @@ -280,7 +281,8 @@ class LightningModule(ABC, GradInformation, ModelIO, ModelHooks): # log training accuracy at the end of an epoch results = { - 'log': {'train_acc': train_acc_mean.item()} + 'log': {'train_acc': train_acc_mean.item()}, + 'progress_bar': {'train_acc': train_acc_mean}, } return results @@ -303,6 +305,7 @@ class LightningModule(ABC, GradInformation, ModelIO, ModelHooks): # log training accuracy at the end of an epoch results = { 'log': {'train_acc': train_acc_mean.item(), 'step': self.current_epoch} + 'progress_bar': {'train_acc': train_acc_mean}, } return results """ diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py index 073403c9..2eb033c0 100644 --- a/pytorch_lightning/trainer/training_loop.py +++ b/pytorch_lightning/trainer/training_loop.py @@ -491,6 +491,7 @@ class TrainerTrainLoopMixin(ABC): callback_epoch_metrics = _processed_outputs[3] self.log_metrics(log_epoch_metrics, {}) self.callback_metrics.update(callback_epoch_metrics) + self.add_progress_bar_metrics(_processed_outputs[1]) # when no val loop is present or fast-dev-run still need to call checkpoints if not self.is_overridden('validation_step') and not (self.fast_dev_run or should_check_val): diff --git a/tests/models/test_module_hooks.py b/tests/models/test_module_hooks.py new file mode 100644 index 00000000..c2f5056b --- /dev/null +++ b/tests/models/test_module_hooks.py @@ -0,0 +1,46 @@ +import torch + +from pytorch_lightning import Trainer +from tests.base import EvalModelTemplate + +import tests.base.utils as tutils + + +def test_training_epoch_end_metrics_collection(tmpdir): + """ Test that progress bar metrics also get collected at the end of an epoch. """ + num_epochs = 3 + class CurrentModel(EvalModelTemplate): + + def training_step(self, *args, **kwargs): + output = super().training_step(*args, **kwargs) + output['progress_bar'].update({'step_metric': torch.tensor(-1)}) + output['progress_bar'].update({'shared_metric': 100}) + return output + + def training_epoch_end(self, outputs): + epoch = self.current_epoch + # both scalar tensors and Python numbers are accepted + return { + 'progress_bar': { + f'epoch_metric_{epoch}': torch.tensor(epoch), # add a new metric key every epoch + 'shared_metric': 111, + } + } + + model = CurrentModel(tutils.get_default_hparams()) + trainer = Trainer( + max_epochs=num_epochs, + default_root_dir=tmpdir, + overfit_pct=0.1, + ) + result = trainer.fit(model) + assert result == 1 + metrics = trainer.progress_bar_dict + + # metrics added in training step should be unchanged by epoch end method + assert metrics['step_metric'] == -1 + # a metric shared in both methods gets overwritten by epoch_end + assert metrics['shared_metric'] == 111 + # metrics are kept after each epoch + for i in range(num_epochs): + assert metrics[f'epoch_metric_{i}'] == i