Lightning-AI / Lightning-AI/pytorch-lightning
Fit loop and validation loop tear down does not dereference passed dataloaders
@gov-ind is already working on this.
Since Aug 12, 2023.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Bug description
When a dataloader is passed to `pl.Trainer`'s `fit`, the trainer does not seem to clear all references to the passed dataloader. As a result, pickling the trainer will pickle the whole dataloader, which isn't ideal for large dataset.
### Steps to reporoduce
Run the MRE below.
### Expected Behaviour
The trainer's size after fitting should not increase by 1MB (the size of the dataset).
### Actual Behaviour
The trainer's size after fitting increased by 1MB (the size of the dataset).
### What version are you seeing the problem on?
master
### How to reproduce the bug
```python
from pickle import dumps
import lightning.pytorch as pl
from lightning.pytorch.demos.boring_classes import BoringModel, RandomDataset
import torch
model = BoringModel()
dataset = RandomDataset(32, 8192)
data_size = len(dumps(dataset))
trainer = pl.Trainer(fast_dev_run=True)
prev_size = len(dumps(trainer))
trainer.fit(model, train_dataloaders=torch.utils.data.DataLoader(dataset))
new_size = len(dumps(trainer))
size_diff_mb = (new_size - prev_size) // (1024 ** 2)
data_size_mb = data_size // (1024 ** 2)
print(size_diff_mb == data_size_mb)
assert size_diff_mb == 0
```
### Error messages and logs
_No response_
### Environment
Current environment
* CUDA:
- GPU: None
- available: False
- version: None
* Lightning:
- lightning: 2.0.6
- lightning-cloud: 0.5.37
- lightning-utilities: 0.9.0
- pytorch-lightning: 2.0.6
- torch: 2.0.1+cpu
- torch-model-archiver: 0.8.1
- torch-workflow-archiver: 0.2.9
- torchmetrics: 0.11.4
- torchserve: 0.8.1
- torchvision: 0.15.2a0
* Packages:
- absl-py: 1.4.0
- anyio: 3.7.1
- arrow: 1.2.3
- astunparse: 1.6.3
- attrs: 23.1.0
- backoff: 2.2.1
- backports.cached-property: 1.0.2
- backports.functools-lru-cache: 1.6.5
- beautifulsoup4: 4.12.2
- blessed: 1.19.1
- blinker: 1.4
- brotlipy: 0.7.0
- build: 0.10.0
- cachecontrol: 0.12.11
- cachetools: 4.2.2
- captum: 0.6.0
- certifi: 2023.7.22
- cffi: 1.15.0
- charset-normalizer: 3.2.0
- cleo: 2.0.1
- click: 8.1.6
- colorama: 0.4.6
- contourpy: 1.1.0
- crashtest: 0.4.1
- croniter: 1.4.1
- cryptography: 41.0.2
- cycler: 0.11.0
- dateutils: 0.6.12
- deepdiff: 5.8.1
- distlib: 0.3.7
- docker: 6.1.3
- docstring-parser: 0.15
- dulwich: 0.21.3
- enum-compat: 0.0.3
- exceptiongroup: 1.1.2
- fastapi: 0.100.1
- filelock: 3.12.2
- flatbuffers: 2.0
- fonttools: 4.42.0
- fsspec: 2023.6.0
- gast: 0.4.0
- ghp-import: 2.1.0
- google-api-core: 2.11.1
- google-auth: 1.21.3
- google-auth-oauthlib: 0.5.2
- google-cloud-core: 2.3.3
- google-cloud-storage: 2.10.0
- google-crc32c: 1.5.0
- google-pasta: 0.2.0
- google-resumable-media: 2.5.0
- googleapis-common-protos: 1.59.1
- griffe: 0.32.3
- grpcio: 1.48.2
- h11: 0.14.0
- h5py: 3.7.0
- html5lib: 1.1
- idna: 3.4
- importlib-metadata: 6.8.0
- importlib-resources: 6.0.0
- iniconfig: 2.0.0
- inquirer: 3.1.3
- installer: 0.7.0
- itsdangerous: 2.1.2
- jaraco.classes: 3.3.0
- jeepney: 0.8.0
- jinja2: 3.1.2
- joblib: 1.3.1
- jsonschema: 4.17.3
- keras: 2.12.0
- keras-pickle-wrapper: 1.0.5
- keras-preprocessing: 1.1.2
- keyring: 23.13.1
- kfp: 2.0.1
- kfp-pipeline-spec: 0.2.2
- kfp-server-api: 2.0.0
- kiwisolver: 1.4.4
- kubernetes: 26.1.0
- lightning: 2.0.6
- lightning-cloud: 0.5.37
- lightning-utilities: 0.9.0
- lockfile: 0.12.2
- markdown: 3.4.1
- markdown-it-py: 3.0.0
- markupsafe: 2.1.1
- material: 0.1
- materialx: 0.0.0.dev1
- matplotlib: 3.7.2
- mdurl: 0.1.0
- mergedeep: 1.3.4
- minio: 7.1.15
- mkdocs: 1.5.2
- mkdocs-autorefs: 0.5.0
- mkdocs-material: 9.1.21
- mkdocs-material-extensions: 1.1.1
- mkdocstrings: 0.22.0
- mkdocstrings-python: 1.3.0
- mlframework: 2.2.1
- more-itertools: 10.0.0
- mpmath: 1.3.0
- msgpack: 1.0.3
- networkx: 3.1
- numpy: 1.22.3
- oauthlib: 3.2.2
- opt-einsum: 3.3.0
- ordered-set: 4.1.0
- packaging: 23.1
- pandas: 2.0.3
- pathspec: 0.11.2
- pexpect: 4.8.0
- pillow: 9.4.0
- pip: 23.2.1
- pipdeptree: 2.12.0
- pkginfo: 1.9.6
- pkgutil-resolve-name: 1.3.10
- platformdirs: 3.10.0
- pluggy: 1.2.0
- poetry: 1.5.1
- poetry-core: 1.6.1
- poetry-plugin-export: 1.4.0
- protobuf: 3.20.3
- psutil: 5.9.0
- psycopg2: 2.9.6
- ptyprocess: 0.7.0
- pyasn1: 0.4.8
- pyasn1-modules: 0.2.8
- pycparser: 2.21
- pydantic: 1.10.8
- pygments: 2.15.1
- pyjwt: 2.8.0
- pymdown-extensions: 10.1
- pympler: 1.0.1
- pyopenssl: 23.2.0
- pyparsing: 3.0.9
- pyproject-hooks: 1.0.0
- pyrsistent: 0.18.0
- pysocks: 1.7.1
- pytest: 7.4.0
- python-dateutil: 2.8.2
- python-editor: 1.0.4
- python-multipart: 0.0.6
- pytorch-lightning: 2.0.6
- pytz: 2023.3
- pyyaml: 6.0.1
- pyyaml-env-tag: 0.1
- rapidfuzz: 2.13.7
- readchar: 4.0.5.dev0
- regex: 2023.8.8
- requests: 2.31.0
- requests-oauthlib: 1.3.0
- requests-toolbelt: 0.9.1
- rich: 13.5.1
- rsa: 4.7.2
- scikit-learn: 1.3.0
- scipy: 1.11.1
- secretstorage: 3.3.3
- setuptools: 68.0.0
- shellingham: 1.5.1
- six: 1.16.0
- sniffio: 1.3.0
- soupsieve: 2.3.2.post1
- starlette: 0.27.0
- starsessions: 1.3.0
- sympy: 1.12
- tabulate: 0.9.0
- tenacity: 8.2.2
- tensorboard: 2.12.1
- tensorboard-data-server: 0.7.0
- tensorboard-plugin-wit: 1.8.1
- tensorflow: 2.12.0
- tensorflow-estimator: 2.12.0
- termcolor: 2.1.0
- threadpoolctl: 3.2.0
- tomli: 2.0.1
- tomlkit: 0.12.1
- torch: 2.0.1+cpu
- torch-model-archiver: 0.8.1
- torch-workflow-archiver: 0.2.9
- torchmetrics: 0.11.4
- torchserve: 0.8.1
- torchvision: 0.15.2a0
- tqdm: 4.65.0
- traitlets: 5.9.0
- trove-classifiers: 2023.7.6
- typeguard: 4.0.0
- typing-extensions: 4.7.1
- tzdata: 2023.3
- urllib3: 1.26.15
- uvicorn: 0.23.2
- virtualenv: 20.24.2
- watchdog: 3.0.0
- wcwidth: 0.2.6
- webencodings: 0.5.1
- websocket-client: 1.6.1
- websockets: 10.4
- werkzeug: 2.2.3
- wheel: 0.38.4
- wrapt: 1.14.1
- zipp: 3.16.2
* System:
- OS: Linux
- architecture:
- 64bit
- ELF
- processor: x86_64
- python: 3.10.12
- release: 5.4.0-1106-gcp
- version: #115~18.04.1-Ubuntu SMP Mon May 22 20:46:39 UTC 202 3
### More info
### Possible Fix
Set `trainer.fit_loop._data_source.instance = None` and `trainer.fit_loop._combined_loader = None` in the loop's `teardown`. These two variables seem to be holding on to the dataloader's reference. I can create a PR that applies this fix.
### Workaround
1. Set the above variables to `None` manually just before pickling the trainer.
2. Since this issue doesn't happen if the dataloader is passed by overriding `train_dataloader`, don't pass the dataloader to `fit`.
cc @justusschock @awaelchli @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.
Assessment
This issue has not been assessed yet.