Repository: Lightning-AI/pytorch-lightning
Pull request: fix collecting training_step outputs
## What does this PR do?

Fixes #8603

With git bisect found that https://github.com/PyTorchLightning/pytorch-lightning/pull/7891/ had introduced the issue by removing a "copy" of the training step output [here](https://github.com/PyTorchLightning/pytorch-lightning/pull/7891#discussion_r679001624).

Tests fail on master, reproducing the issue. 

### Does your PR introduce any breaking changes? If yes, please list them.

<!-- FILL IN or None -->

## Before submitting
- [x] Was this discussed/approved via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your PR does only one thing, instead of bundling different changes together?
- [x] Did you make sure to update the documentation with your changes? (if necessary)
- [x] Did you write any new necessary tests? (not for typos and docs)
- [x] Did you verify new and existing tests pass locally with your changes?
- [x] Did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- For CHANGELOG separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review
Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

 - [x] Is this pull request ready for review? (if not, please submit in draft mode)
 - [x] Check that all items from **Before submitting** are resolved
 - [x] Make sure the title is self-explanatory and the description concisely explains the PR
 - [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?

not this time 😢 

Problem statement:
outputs in training_epoch_end contain only outputs from last batch repeated
## 🐛 Bug

outputs in training_epoch_end contain only outputs from last batch repeated multiple times. I believe it got broken only in 1.4.0, but in 1.3.x it worked.

### To Reproduce

```python
import torch
from torch.utils.data import DataLoader, Dataset
from pytorch_lightning import LightningModule, Trainer


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        print(f'training_step, {batch_idx=}: {loss=}')
        return loss

    def configure_optimizers(self):
        return torch.optim.SGD(self.parameters(), lr=0.1)

    def training_epoch_end(self, outputs):
        print('training_epoch_end:', outputs)


dl = DataLoader(RandomDataset(32, 100), batch_size=10)

model = BoringModel()
trainer = Trainer(max_epochs=1, progress_bar_refresh_rate=0)
trainer.fit(model, dl)
```

This will print the same loss repeated 10 times (equal to last batch loss) in training_epoch_end:
```
training_step, batch_idx=0: loss=tensor(0.6952, grad_fn=<SumBackward0>)
training_step, batch_idx=1: loss=tensor(-18.9661, grad_fn=<SumBackward0>)
training_step, batch_idx=2: loss=tensor(-27.7834, grad_fn=<SumBackward0>)
training_step, batch_idx=3: loss=tensor(-84.3158, grad_fn=<SumBackward0>)
training_step, batch_idx=4: loss=tensor(-119.3664, grad_fn=<SumBackward0>)
training_step, batch_idx=5: loss=tensor(-138.1930, grad_fn=<SumBackward0>)
training_step, batch_idx=6: loss=tensor(-126.4004, grad_fn=<SumBackward0>)
training_step, batch_idx=7: loss=tensor(-143.7022, grad_fn=<SumBackward0>)
training_step, batch_idx=8: loss=tensor(-175.9583, grad_fn=<SumBackward0>)
training_step, batch_idx=9: loss=tensor(-161.6977, grad_fn=<SumBackward0>)

training_epoch_end: [{'loss': tensor(-161.6977)}, {'loss': tensor(-161.6977)}, {'loss': tensor(-161.6977)}, {'loss': tensor(-161.6977)}, {'loss': tensor(-161.6977)}, {'loss': tensor(-161.6977)}, {'loss': tensor(-161.6977)}, {'loss': tensor(-161.6977)}, {'loss': tensor(-161.6977)}, {'loss': tensor(-161.6977)}]
```


### Expected behavior

Output from all steps/batches is available in training_epoch_end (not only from last batch)

### Environment

```
* CUDA:
	- GPU:
	- available:         False
	- version:           None
* Packages:
	- numpy:             1.18.5
	- pyTorch_debug:     False
	- pyTorch_version:   1.8.0
	- pytorch-lightning: 1.4.0
	- tqdm:              4.47.0
* System:
	- OS:                Darwin
	- architecture:
		- 64bit
		-
	- processor:         i386
	- python:            3.8.3
	- version:           Darwin Kernel Version 19.6.0: Tue Jan 12 22:13:05 PST 2021; root:xnu-6153.141.16~1/RELEASE_X86_64
```

Hints from issue discussion:

Commit message: add a comment for gc.collect()

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index a33cdd17031bd..5c320befec7b6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -71,7 +71,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
--
+- Fixed an issue with `training_step` outputs not getting collected correctly for `training_epoch_end` ([#8613](https://github.com/PyTorchLightning/pytorch-lightning/pull/8613))
 
 
 -
diff --git a/pytorch_lightning/loops/batch/training_batch_loop.py b/pytorch_lightning/loops/batch/training_batch_loop.py
index 3b49fe4a1f63a..4859b1060c7c8 100644
--- a/pytorch_lightning/loops/batch/training_batch_loop.py
+++ b/pytorch_lightning/loops/batch/training_batch_loop.py
@@ -14,6 +14,7 @@
 
 from collections import OrderedDict
 from contextlib import contextmanager
+from copy import copy
 from functools import partial, update_wrapper
 from typing import Any, Callable, Dict, Generator, List, Mapping, Optional, Tuple
 
@@ -143,12 +144,12 @@ def advance(self, batch, batch_idx, dataloader_idx):
 
                 result = self._run_optimization(batch_idx, split_batch, opt_idx, optimizer)
                 if result:
-                    self.batch_outputs[opt_idx].append(result.training_step_output)
+                    self.batch_outputs[opt_idx].append(copy(result.training_step_output))
         else:
             # in manual optimization, there is no looping over optimizers
             result = self._run_optimization(batch_idx, split_batch)
             if result:
-                self.batch_outputs[0].append(result.training_step_output)
+                self.batch_outputs[0].append(copy(result.training_step_output))
 
     def teardown(self) -> None:
         # release memory
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
index 212086dc0f678..513ea46f4695c 100644
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Trainer to automate the training."""
+import gc
 import logging
 import os
 import traceback
@@ -982,6 +983,8 @@ def _post_dispatch(self):
         self.accelerator.teardown()
         self._active_loop.teardown()
         self.logger_connector.teardown()
+        # force release any leftover memory, including CUDA tensors
+        gc.collect()
 
     def _dispatch(self):
         if self.evaluating:
diff --git a/tests/trainer/loops/test_training_loop_flow_dict.py b/tests/trainer/loops/test_training_loop_flow_dict.py
index f064dacb78844..ab4d7979bbf39 100644
--- a/tests/trainer/loops/test_training_loop_flow_dict.py
+++ b/tests/trainer/loops/test_training_loop_flow_dict.py
@@ -108,7 +108,7 @@ def training_step(self, batch, batch_idx):
             acc = acc + batch_idx
 
             self.training_step_called = True
-            out = {"loss": acc, "random_things": [1, "a", torch.tensor(2)]}
+            out = {"loss": acc, "random_things": [1, "a", torch.tensor(2)], "batch_idx": batch_idx}
             return out
 
         def training_epoch_end(self, outputs):
@@ -116,11 +116,13 @@ def training_epoch_end(self, outputs):
 
             # verify we saw the current num of batches
             assert len(outputs) == 2
+            assert len({id(output) for output in outputs}) == 2
+            assert [output["batch_idx"] for output in outputs] == [0, 1]
 
             for b in outputs:
                 assert isinstance(b, dict)
                 assert self.count_num_graphs(b) == 0
-                assert {"random_things", "loss"} == set(b.keys())
+                assert {"random_things", "loss", "batch_idx"} == set(b.keys())
 
         def backward(self, loss, optimizer, optimizer_idx):
             return LightningModule.backward(self, loss, optimizer, optimizer_idx)
@@ -155,7 +157,7 @@ def training_step(self, batch, batch_idx):
             acc = acc + batch_idx
 
             self.training_step_called = True
-            self.out = {"loss": acc, "random_things": [1, "a", torch.tensor(2)]}
+            self.out = {"loss": acc, "random_things": [1, "a", torch.tensor(2)], "batch_idx": batch_idx}
             return self.out
 
         def training_step_end(self, tr_step_output):
@@ -169,11 +171,13 @@ def training_epoch_end(self, outputs):
 
             # verify we saw the current num of batches
             assert len(outputs) == 2
+            assert len({id(output) for output in outputs}) == 2
+            assert [output["batch_idx"] for output in outputs] == [0, 1]
 
             for b in outputs:
                 assert isinstance(b, dict)
                 assert self.count_num_graphs(b) == 0
-                assert {"random_things", "loss"} == set(b.keys())
+                assert {"random_things", "loss", "batch_idx"} == set(b.keys())
 
         def backward(self, loss, optimizer, optimizer_idx):
             return LightningModule.backward(self, loss, optimizer, optimizer_idx)


Code review comments:
diff_hunk: @@ -982,6 +983,8 @@ def _post_dispatch(self):
         self.accelerator.teardown()
         self._active_loop.teardown()
         self.logger_connector.teardown()
+        # force release any leftover memory, including CUDA tensors
+        gc.collect()
path: pytorch_lightning/trainer/trainer.py
text: @user1:
what kind of memory is this specifically? So where is it coming from? If you just call the gc without explicitly freeing it anywhere (e.g by resetting variables, decreasing counters or marking them as safe to delete with `del`), we should be fine with skipping it and having the gc release the memory if it's needed somewhere, no?

@author:
I agree but we have [this](https://github.com/PyTorchLightning/pytorch-lightning/blob/c99e2fe0d2bf713f35054eaa0d521ee7f6030786/tests/trainer/test_trainer.py#L1895) test here which measures the memory right after the fit call. 
The test was flaky before, and it was written when this bug in this PR existed so now that we fixed it we use more memory and gc is not quick enough to release everything right after the fit() call. 

Would you rather put it into the test?

@user1:
I think this is a more general question. IMO it comes down to whether we want to explicitly free anything or just rely on python to do so and just make sure, we reset everything we don't need anymore.

In the first case, I think we should do that more aggressive across the whole codebase. In the second case, I think we should move this call to the test, since it then is only relevant for it (and at least for stuff that can be freed like that we don't hold a reference anymore which is fine from our side then)

@author:
Requesting quick comment from @user2 and @user3 for their opinion. Both is fine with me.

The test was introduced here #8490 and the motivation for clearing the memory completely after fit was to prevent oom issues when running many fit() in sequence.

Option A: gc.collect() in our trainer teardown
Option B: gc.collect() directly in our test

@user2:
I would go with `Option B` for now.