Repository: Lightning-AI/pytorch-lightning

Issue:
Log training metrics for each epoch
Currently, I am able to log training metrics to Tensorboard using:

```
import pytorch_lightning as pl
from pytorch_lightning.loggers import TensorBoardLogger

logger = TensorBoardLogger(save_dir=save_dir, name="my_model")

[...]

trainer = pl.Trainer(logger=logger)
```

This logs training metrics (loss, for instance) after each batch. I would like to be able to average these metrics across all batches and log them to TensorBoard only once, at the end of each epoch. This is what the `validation_end` method does in your example: https://github.com/PyTorchLightning/pytorch-lightning/blob/446a1e23d7fe3b2e07f1a5887fe819d0dfa7d4e0/pl_examples/basic_examples/lightning_module_template.py#L145. 

I first thought about writing my own `training_end` method. But this method is called after each batch instead of being called at the end of an epoch (as I would have thought). The method `on_epoch_end` seems interesting but does not receive an `outputs` argument as `training_end` does. Basically, in my model, I would like to write something like:  `self.logger.experiment.add_scalar('training_loss', train_loss_mean, global_step=self.current_epoch)`, but I do not know where to put this line.

 - OS: Debian GNU/Linux 9.11 (stretch)
 - Packaging: PIP
 - Version 0.6.1.dev0



Discussion:
Hey, thanks for your contribution! Great first issue!
How about this:
In `__init__`: 
`self.training_losses = []`
In training_step method:
`self.training_losses.append(loss.item())`
In epoch_end method:
```
train_loss_mean = np.mean(self.training_losses)
self.logger.experiment.add_scalar('training_loss', train_loss_mean, global_step=self.current_epoch)
self.training_losses = []  # reset for next epoch
```
I'm also trying to log training metrics at the end of each epoch, and tried it as follows (based off the MNIST example in https://colab.research.google.com/drive/1F_RNcHzTfFuQf-LeKvSlud6x7jXYkG31#scrollTo=x-34xKCI40yW):

```
class CoolSystem(pl.LightningModule):

    def __init__(self, hparams):
        super(CoolSystem, self).__init__()

        self.l1 = torch.nn.Linear(28 * 28, 10)
        
    # ...

    def training_step(self, batch, batch_idx):
        # REQUIRED
        x, y = batch
        y_hat = self.forward(x)
        
        loss = F.cross_entropy(y_hat, y)
        
        _, preds = torch.max(y_hat, dim=1)
        acc = torch.sum(preds == y.data) / (y.shape[0] * 1.0)

        return {'train_loss': loss, 'train_acc': acc}
    
    def training_end(self, outputs):
        # OPTIONAL
        avg_loss = torch.stack([x['train_loss'] for x in outputs]).mean() # -------- Error is thrown here
        avg_acc = torch.stack([x['train_acc'].float() for x in outputs]).mean()

        logs = {'train_loss': avg_loss, 'train_acc': avg_acc}
        return {'avg_train_loss': avg_loss, 'avg_train_acc': avg_acc, 'log': logs, 'progress_bar': logs}

    def validation_step(self, batch, batch_idx):
        # OPTIONAL
        x, y = batch
        y_hat = self.forward(x)

        _, preds = torch.max(y_hat, 1)
        acc = torch.sum(preds == y.data) / (y.shape[0] * 1.0)

        return {'val_loss': F.cross_entropy(y_hat, y), 'val_acc': acc}

    def validation_end(self, outputs):
        # OPTIONAL
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        avg_acc = torch.stack([x['val_acc'].float() for x in outputs]).mean()

        print (avg_acc)

        logs = {'val_loss': avg_loss, 'val_acc': avg_acc}
        return {'avg_val_loss': avg_loss, 'avg_val_acc': avg_acc, 'log': logs, 'progress_bar': logs}
```

However, the following error is thrown at `training_end`:
```
TypeError: string indices must be integers
```

Weird thing is that it works for `validation_end` (the validation check passes, and no error is thrown when I remove `training_end` from my system to debug); could you pls advise? Thanks!

@awaelchli That works for me, thanks!

@polars05 From my understanding, `training_step`  and `validation_step` compute the metrics for the current step (batch) and return them in the `output` dict. In the validation loop, the metrics from all steps (all batches in an epoch) are aggregated in `outputs`. The method `validation_end` allows you to average those metrics epoch-wise (for instance, compute a single loss score for the whole validation dataset). What was confusing for me is that `training_end` does not work like `validation_end` (see my post). @awaelchli feel free to correct me if I'm wrong!
yeah, training_end is a bit confusing. we discussed this and decided to change the name/modify it. 

training_end aggregates outputs of a batch on dp. We need a true training_end and somethinf for the dp aggregation. 

@jeremyjordan @ethanwharris i think you guys had opinions about this? should we try to sneak this into this next release?
@jbschiratti, thanks for the clarifications! I've modified my system as per what @awaelchli suggested to you and that works for me as well.

However, in the process, I observed that under `training_step`, if we did `return {'train_loss': loss} ` instead of `return {'loss': loss}`, then the following error is thrown:
```
RuntimeError: No `loss` value in the dictionary returned from `model.training_step()`.
```

For reference, here's my system:

```
class CoolSystem(pl.LightningModule):

    def __init__(self, hparams):
        super(CoolSystem, self).__init__()

        self.hparams = hparams
        self.data_dir = self.hparams.data_dir
        # self.classes = self.hparams.classes
        self.num_classes = self.hparams.num_classes # len(self.hparams.classes)

        self.model = torchvision.models.resnet18(pretrained=True) # final layer is of size [bs, 1000]
        num_ftrs = self.model.fc.in_features
        self.model.fc = torch.nn.Linear(num_ftrs, self.num_classes) # change final layer to be of size [bs, 2]
        
        self.training_acc_across_batches_at_curr_epoch = []

    # ...

    def training_step(self, batch, batch_idx):
        # REQUIRED
        x, y = batch
        y_hat = self.forward(x)
        
        loss = F.cross_entropy(y_hat, y)
        
        _, preds = torch.max(y_hat, dim=1)
        acc = torch.sum(preds == y.data) / (y.shape[0] * 1.0)
        # torch.tensor(acc)
        self.training_acc_across_batches_at_curr_epoch.append(acc.item())
        
        # return {'loss': loss}         # -------> this works...
        return {'train_loss': loss}   # --------> ... but this doesn't
    
    def on_epoch_end(self):
        train_acc_mean = np.mean(self.training_acc_across_batches_at_curr_epoch)
        
        self.logger.experiment.add_scalar('train_acc', train_acc_mean, global_step=self.current_epoch)
        self.training_acc_across_batches_per_epoch = []  # reset for next epoch
```

**My guess is that we do not have the flexibility to name the dict key in `training_step`?**

In addition, if we did the following instead in `training_step`:
```
logs = {'train_loss': loss, 'train_acc': acc}
return {'loss': loss, 'train_acc': acc, 'log': logs, 'progress_bar': logs}
```
then both `loss` and `train_loss` shows up on the progress bar

I could simply do:
```
dict_for_progress_bar = {'train_acc': acc}
dict_for_log = {'train_loss': loss, 'train_acc': acc}
return {'loss': loss, 'train_acc': acc, 'log': dict_for_logs, 'progress_bar': dict_for_progress_bar}
```
**but it might be cleaner for the user to not have to create two separate dicts?**
@williamFalcon  I agree that the behavior for `training_end` should be symmetric with `validation_end`. Unfortunately, this does require changing the API - should we include a warning message in the next release to alert of a change upcoming in the subsequent release? 

The current `training_end` could be renamed to something more explicit such as `collect_dp_batches` or something of the like. 
@polars05 that is correct! you can see how the values returned are processed here 
https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/trainer/logging.py#L90

on line 153 you can see that there's an explicit check for a `loss` key in the output dict.
Just wondering what is the best way to do something at the end of each training epoch, before validation starts?
From this thread, I assume training_end is not the right place to do that?
Currently, I had to do something in validation_step, where I check whether the index == 0 and do something, but this seems very ugly.
@versatran01 Maybe for your use case the on_epoch_end() hook is enough? It doesn't get the metrics from training though. 

If you need something that is conceptually similar as validation_end for training, then this does not exist yet. And this is also what @jbschiratti is asking for in this issue.
If you need a workaround until the changes are implemented, see my answer above.
@awaelchli Thanks for the suggestion. I just saw some huge updates to the callback system, will give it a try.
@williamFalcon @jeremyjordan Changing the current `training_end` method into something like `collect_dp_batches` and making `training_end` similar to the current `validation_end` does not seem to be too much work. Shall I submit a PR for this?
training_end makes sense.

maybe a better name for collect_dp_batches is: ```training_step_reduce``` or something like that
updated in 0.7.1
@williamFalcon I had a look at the latest release and I think that my issue still stands. Unless I am mistaken, there is still no equivalent of `validation_end` (now called `validation_epoch_end`) in [training_loop.py](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/trainer/training_loop.py). 

At each epoch, one can average validation loss across all batches by doing:
```
def validation_epoch_end(self, outputs):
    val_loss_mean = 0
    for output in outputs:
        val_loss = output['val_loss']
        if self.trainer.use_dp or self.trainer.use_ddp2:
            val_loss = torch.mean(val_loss)
        val_loss_mean += val_loss
    val_loss_mean /= len(outputs)
    tqdm_dict = {'val_loss': val_loss_mean}
    return OrderedDict({'val_loss': val_loss_mean, 'progress_bar': tqdm_dict, 'log': tqdm_dict})
```
But as of today, there is still no way to do the same in the training loop ! I though about the function `training_epoch_end` but it is not used. It is only mentioned in a comment on [line 719](https://github.com/PyTorchLightning/pytorch-lightning/blob/ee68d5ba8e66bba52b3cc3d9d8780d6df2f3204b/pytorch_lightning/trainer/training_loop.py#L718).
Shall I open a new issue ?
@jbschiratti you're right, i don't see where this is implemented. reopened the issue.

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
@@ -231,6 +231,78 @@ def training_end(self, *args, **kwargs):
             Deprecated in v0.7.0. use training_step_end instead
         """
 
+    def training_epoch_end(
+            self,
+            outputs: Union[List[Dict[str, Tensor]], List[List[Dict[str, Tensor]]]]
+    ) -> Dict[str, Dict[str, Tensor]]:
+        """Called at the end of training epoch with the outputs of all training_steps
+
+        .. code-block:: python
+
+            # the pseudocode for these calls
+
+            train_outs = []
+            for train_batch in train_data:
+                out = training_step(train_batch)
+                train_outs.append(out)
+            training_epoch_end(val_outs)
+
+        Args:
+            outputs: List of outputs you defined in training_step, or if there are multiple
+            dataloaders, a list containing a list of outputs for each dataloader
+
+        Return:
+            Dict or OrderedDict (dict): Dict has the following optional keys:
+            progress_bar -> Dict for progress bar display. Must have only tensors
+            log -> Dict of metrics to add to logger. Must have only tensors (no images, etc)
+
+        .. note:: If this method is not overridden, this won't be called.
+
+        - The outputs here are strictly for logging or progress bar.
+        - If you don't need to display anything, don't return anything.
+        - If you want to manually set current step, you can specify the 'step' key in the 'log' Dict
+
+        Examples:
+            With a single dataloader
+
+            .. code-block:: python
+
+                def training_epoch_end(self, outputs):
+                    train_acc_mean = 0
+                    for output in outputs:
+                        train_acc_mean += output['train_acc']
+
+                    train_acc_mean /= len(outputs)
+
+                    # log training accuracy at the end of an epoch
+                    results = {
+                        'log': {'train_acc': train_acc_mean.item()}
+                    }
+                    return results
+
+            With multiple dataloaders, `outputs` will be a list of lists. The outer list contains
+            one entry per dataloader, while the inner list contains the individual outputs of
+            each validation step for that dataloader.
+
+            .. code-block:: python
+
+                def training_epoch_end(self, outputs):
+                    train_acc_mean = 0
+                    i = 0
+                    for dataloader_outputs in outputs:
+                        for output in dataloader_outputs:
+                            train_acc_mean += output['train_acc']
+                            i += 1
+
+                    train_acc_mean /= i
+
+                    # log training accuracy at the end of an epoch
+                    results = {
+                        'log': {'train_acc': train_acc_mean.item(), 'step': self.current_epoch}
+                    }
+                    return results
+        """
+
     def training_step_end(self, *args, **kwargs) -> Dict[
         str, Union[Tensor, Dict[str, Tensor]]
     ]:
@@ -453,7 +525,7 @@ def validation_epoch_end(
             outputs: Union[List[Dict[str, Tensor]], List[List[Dict[str, Tensor]]]]
     ) -> Dict[str, Dict[str, Tensor]]:
         """
-        Called at end of validation epoch with the output of all validation_steps
+        Called at end of validation epoch with the outputs of all validation_steps
 
         .. code-block:: python
 
@@ -462,7 +534,7 @@ def validation_epoch_end(
             val_outs = []
             for val_batch in val_data:
                 out = validation_step(train_batch)
-                train_outs.append(out)
+                val_outs.append(out)
             validation_epoch_end(val_outs)
 
         Args:
@@ -493,7 +565,7 @@ def validation_epoch_end(self, outputs):
                     val_acc_mean /= len(outputs)
                     tqdm_dict = {'val_acc': val_acc_mean.item()}
 
-                    # show val_loss and val_acc in progress bar but only log val_loss
+                    # show val_acc in progress bar but only log val_loss
                     results = {
                         'progress_bar': tqdm_dict,
                         'log': {'val_acc': val_acc_mean.item()}
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
@@ -145,6 +145,7 @@ def training_step(self, batch, batch_idx):
 from pytorch_lightning.callbacks.base import Callback
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.loggers import LightningLoggerBase
+from pytorch_lightning.overrides.data_parallel import LightningDistributedDataParallel, LightningDataParallel
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.trainer.supporters import TensorRunningMean
 
@@ -390,6 +391,9 @@ def train(self):
 
     def run_training_epoch(self):
 
+        # get model
+        model = self.get_model()
+
         # Epoch start events
         with self.profiler.profile('on_epoch_start'):
             # callbacks
@@ -397,7 +401,7 @@ def run_training_epoch(self):
 
             # model hooks
             if self.is_function_implemented('on_epoch_start'):
-                self.get_model().on_epoch_start()
+                model.on_epoch_start()
 
         # track local dataloader so TPU can wrap each epoch
         train_dataloader = self.train_dataloader
@@ -408,6 +412,9 @@ def run_training_epoch(self):
             train_dataloader = xla_pl.ParallelLoader(train_dataloader, [device])
             train_dataloader = train_dataloader.per_device_loader(device)
 
+        # bookkeeping
+        outputs = []
+
         # run epoch
         for batch_idx, (batch, is_last_batch) in self.profiler.profile_iterable(
             enumerate(_with_is_last(train_dataloader)), "get_train_batch"
@@ -418,14 +425,15 @@ def run_training_epoch(self):
 
             self.batch_idx = batch_idx
 
-            model = self.get_model()
             model.global_step = self.global_step
 
             # ---------------
             # RUN TRAIN STEP
             # ---------------
-            output = self.run_training_batch(batch, batch_idx)
-            batch_result, grad_norm_dic, batch_step_metrics = output
+            _outputs = self.run_training_batch(batch, batch_idx)
+            batch_result, grad_norm_dic, batch_step_metrics, batch_output = _outputs
+            # detach tensors in batch_output before appending to outputs
+            outputs.append(_recursive_detach(batch_output))
 
             # when returning -1 from train_step, we end epoch early
             early_stop_epoch = batch_result == -1
@@ -484,6 +492,18 @@ def run_training_epoch(self):
             if early_stop_epoch or self.fast_dev_run:
                 break
 
+        # process epoch outputs
+        if isinstance(model, (LightningDistributedDataParallel, LightningDataParallel)):
+            model = model.module
+
+        if self.is_overriden('training_epoch_end', model=model):
+            epoch_output = model.training_epoch_end(outputs)
+            _processed_outputs = self.process_output(epoch_output)
+            log_epoch_metrics = _processed_outputs[2]
+            callback_epoch_metrics = _processed_outputs[3]
+            self.log_metrics(log_epoch_metrics, {})
+            self.callback_metrics.update(callback_epoch_metrics)
+
         # in case validation step is missing and you are not running fast-dev to duplicate last batch
         if not self.is_overriden('validation_step') and not (self.fast_dev_run or should_check_val):
             self.call_checkpoint_callback()
@@ -497,7 +517,7 @@ def run_training_epoch(self):
             self.on_epoch_end()
             # model hooks
             if self.is_function_implemented('on_epoch_end'):
-                self.get_model().on_epoch_end()
+                model.on_epoch_end()
 
     def run_training_batch(self, batch, batch_idx):
         # track grad norms
@@ -546,14 +566,13 @@ def run_training_batch(self, batch, batch_idx):
                 def optimizer_closure():
                     # forward pass
                     with self.profiler.profile('model_forward'):
-                        output = self.training_forward(
+                        output_dict = self.training_forward(
                             split_batch, batch_idx, opt_idx, self.hiddens)
 
-                    closure_loss = output[0]
-                    progress_bar_metrics = output[1]
-                    log_metrics = output[2]
-                    callback_metrics = output[3]
-                    self.hiddens = output[4]
+                        # format and reduce outputs accordingly
+                        processed_output = self.process_output(output_dict, train=True)
+
+                    closure_loss, progress_bar_metrics, log_metrics, callback_metrics, self.hiddens = processed_output
 
                     # accumulate loss
                     # (if accumulate_grad_batches = 1 no effect)
@@ -577,10 +596,10 @@ def optimizer_closure():
                         with self.profiler.profile('on_after_backward'):
                             model_ref.on_after_backward()
 
-                    return closure_loss
+                    return closure_loss, output_dict
 
                 # calculate loss
-                loss = optimizer_closure()
+                loss, batch_output = optimizer_closure()
 
                 # check if loss or model weights are nan
                 self.detect_nan_tensors(loss)
@@ -606,7 +625,8 @@ def optimizer_closure():
                     model = self.get_model()
                     with self.profiler.profile('optimizer_step'):
                         model.optimizer_step(self.current_epoch, batch_idx,
-                                             optimizer, opt_idx, optimizer_closure)
+                                             optimizer, opt_idx,
+                                             lambda: optimizer_closure()[0])
 
                     # calculate running loss for display
                     self.running_loss.append(self.batch_loss_value.mean())
@@ -633,7 +653,7 @@ def optimizer_closure():
         # track all metrics for callbacks
         self.callback_metrics.update({k: v for d in all_callback_metrics for k, v in d.items()})
 
-        return 0, grad_norm_dic, all_log_metrics
+        return 0, grad_norm_dic, all_log_metrics, batch_output
 
     def _get_optimizers_iterable(self):
         if not self.optimizer_frequencies:
@@ -732,9 +752,6 @@ def training_forward(self, batch, batch_idx, opt_idx, hiddens):
             warnings.warn('`training_end` was deprecated in 0.7.0 and will be removed 1.0.0.'
                           ' Use training_epoch_end instead', DeprecationWarning)
 
-        # format and reduce outputs accordingly
-        output = self.process_output(output, train=True)
-
         return output
 
     def update_learning_rates(self, interval: str):
@@ -784,3 +801,29 @@ def _with_is_last(iterable):
         last = val
     # yield last, no longer has next
     yield last, True
+
+
+def _recursive_detach(in_dict):
+    """Detach all tensors in `in_dict`.
+
+    May operate recursively if some of the values in `in_dict` are dictionaries
+    which contain instances of `torch.Tensor`. Other types in `in_dict` are
+    not affected by this utility function.
+
+    Parameters
+    ----------
+    in_dict : dict
+
+    Returns
+    -------
+    out_dict : dict
+    """
+    out_dict = {}
+    for k, v in in_dict.items():
+        if isinstance(v, dict):
+            out_dict.update({k: _recursive_detach(v)})
+        elif callable(getattr(v, 'detach', None)):
+            out_dict.update({k: v.detach()})
+        else:
+            out_dict.update({k: v})
+        return out_dict


Test patch:
diff --git a/tests/loggers/test_base.py b/tests/loggers/test_base.py
--- a/tests/loggers/test_base.py
+++ b/tests/loggers/test_base.py
@@ -124,11 +124,16 @@ def test_multiple_loggers_pickle(tmpdir):
 def test_adding_step_key(tmpdir):
     logged_step = 0
 
-    def _validation_end(outputs):
+    def _validation_epoch_end(outputs):
         nonlocal logged_step
         logged_step += 1
         return {"log": {"step": logged_step, "val_acc": logged_step / 10}}
 
+    def _training_epoch_end(outputs):
+        nonlocal logged_step
+        logged_step += 1
+        return {"log": {"step": logged_step, "train_acc": logged_step / 10}}
+
     def _log_metrics_decorator(log_metrics_fn):
         def decorated(metrics, step):
             if "val_acc" in metrics:
@@ -138,7 +143,8 @@ def decorated(metrics, step):
         return decorated
 
     model, hparams = tutils.get_default_model()
-    model.validation_epoch_end = _validation_end
+    model.validation_epoch_end = _validation_epoch_end
+    model.training_epoch_end = _training_epoch_end
     trainer_options = dict(
         max_epochs=4,
         default_save_path=tmpdir,