Lightning-AI / Lightning-AI/pytorch-lightning

CPU-Memory keeps accumulating during `trainer.predict`

Open
#19,398 6 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug performance trainer: predict ver: 2.1.x
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.