Repository: Lightning-AI/pytorch-lightning

Issue:
Also update progress_bar in training_epoch_end
## 🚀 Feature
[PR 1357](https://github.com/PyTorchLightning/pytorch-lightning/pull/1357) implements training_epoch_end to log metrics. The comments in the issue suggest that it should behave like validation_epoch_end, but the PR only replicates the callbacks and metric logging. This feature would add updates to the progress bar as well.

### Motivation
Motivation is for the same usecase as having a progress bar. 
While tensorboard or any other logger is a great resource to have. It is painful to set up in environments like slurm where you can only submit jobs and possibly cannot forward ports or set up the logging.
### Pitch

The changes should only entail adding an update to tqdm [here](https://github.com/PyTorchLightning/pytorch-lightning/blob/b2707c9b2ebeac03f19a3939df9432ac8859d894/pytorch_lightning/trainer/training_loop.py#L503)



Discussion:
Hi! thanks for your contribution!, great first issue!
I will look into this after #1450 is done :)

Resolving patch:
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -257,6 +257,7 @@ def training_epoch_end(
             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 @@ def training_epoch_end(self, outputs):
 
                     # 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 @@ def training_epoch_end(self, outputs):
                     # 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
--- a/pytorch_lightning/trainer/training_loop.py
+++ b/pytorch_lightning/trainer/training_loop.py
@@ -491,6 +491,7 @@ def run_training_epoch(self):
             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_overriden('validation_step') and not (self.fast_dev_run or should_check_val):


Test patch:
diff --git a/tests/models/test_module_hooks.py b/tests/models/test_module_hooks.py
new file mode 100644
--- /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