Lightning-AI / Lightning-AI/pytorch-lightning
CPU-Memory keeps accumulating during `trainer.predict`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Bug description
This is very similar to closed issue #15656
I am working on predicting using PL Trainer on 3D images and these are huge, my process keeps getting killed when a large number of samples are to be predicted. I found #15656 and expected that to be the solution but setting `return_predictions=False` does not fix the memory accumulation.
What seems to work instead is adding a `gc.collect()` in the `predict_loop`. This keeps CPU memory usage constant as would be expected.
It seems like setting `return_predictions=False` should stop the memory accumulation but I'm confused as to why the `gc.collect()` is needed.
This is where the `gc.collect()` is applied: https://github.com/project-lighter/lighter/blob/07018bb2c66c0c8848bab748299e2c2d21c7d185/lighter/callbacks/writer/base.py#L120
I've also attached a memory log using `scalene` of the return predictions and the gc collect comparison. As you can see, there is no memory growth for gc collect.
Would you be able to provide any intuition on this? It would be much appreciated!
### What version are you seeing the problem on?
v2.1
### How to reproduce the bug
_No response_
### Error messages and logs
[gc_collect.pdf](https://github.com/Lightning-AI/pytorch-lightning/files/14146453/gc_collect.pdf)
[return_predictions_false.pdf](https://github.com/Lightning-AI/pytorch-lightning/files/14146454/return_predictions_false.pdf)
### Environment
Current environment
```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):
```
### More info
_No response_
cc @lantiga @borda
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
The issue points to Trainer's predict_loop and to lighter/callbacks/writer/base.py at line 120; start by comparing behavior with return_predictions=False and gc.collect(). No reproduction or Lightning source test is provided, so establish a minimal reproducer before deciding what done means: stable CPU memory without manual collection.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100