Lightning-AI / Lightning-AI/pytorch-lightning

`trainer.fit_loop.setup_data()` does not refresh train dataset in `LightningModule`

Open
#17,327 6 comments 2 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

data handling feature loops pl ver: 2.0.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

PR #16726 replaces the `reset_*_dataloader()` method calls with the respective `Loop.setup_data()` calls. This is also mentioned in the [migration guide](https://lightning.ai/docs/pytorch/stable/upgrade/from_1_9.html#developer).

However, on versions `<= 1.9`, calling `reset_train_dataloader()` would reinstantiate the dataloader from a `LightningModule`'s `train_dataloader()` method. This behaviour is now gone.

My specific use case is that I need to update the dataset of my model during training. I then use `on_train_epoch_end()` or a similar hook to call `reset_train_dataloader()`, to have the updated dataset in the next training epoch. I posted a minimal example below. You can run this example on both `v1.9` and `v2.0` to see the exact difference. `v1.9` runs without problems, whereas `v2.0` fails the second assertion in `training_step()`. I tested it on a fresh conda env install of both versions using python 3.10.

In case I am using the wrong `loop` to call `setup_data()` or am using the new interface incorrectly, please let me know. In that case I would also recommend providing some more hints in the migration guide or on PR #16726 since the current advice is not exactly clear. (i.e. which loops are "top level"?)

### What version are you seeing the problem on?

2.0+

### How to reproduce the bug

```python
try:
import lightning
except ModuleNotFoundError:
import pytorch_lightning as lightning

import torch
from torch.utils.data import DataLoader, TensorDataset

class Model(lightning.LightningModule):

def __init__(self):
super().__init__()
self.train_data = TensorDataset(torch.zeros(1, 1))

def configure_optimizers(self):
return None

def on_train_epoch_end(self):
self.train_data = TensorDataset(torch.ones(1, 1))

if int(lightning.__version__[0]) < 2:
# for version < 2.0 (works)
self.trainer.reset_train_dataloader()
else:
# for version >= 2.0 (does not work)
self.trainer.fit_loop.setup_data()

def train_dataloader(self):
return DataLoader(
self.train_data,
)

def training_step(self, batch, batch_idx):
# de-tuple
batch = batch[0]

if self.trainer.global_step == 0:
assert torch.allclose(batch, torch.zeros_like(batch))
else:
# this assertion fails on lightning v2.0
assert torch.allclose(batch, torch.ones_like(batch))

return torch.tensor(0.0, requires_grad=True)

model = Model()
trainer = lightning.Trainer(max_steps=2)
trainer.fit(model)
```

### Error messages and logs

```python
File "/home/lars/code/python/lightning-trainable/playground.py", line 38, in training_step
assert torch.allclose(batch, torch.ones_like(batch))
AssertionError
```

### Environment

Current environment

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow): Trainer, FitLoop, LightningModule
#- PyTorch Lightning Version (e.g., 1.5.0): 1.9 / 2.0
#- Lightning App Version (e.g., 0.5.2): -
#- PyTorch Version (e.g., 2.0): 1.9 / 2.0
#- Python version (e.g., 3.9): 3.10
#- OS (e.g., Linux): Ubuntu
#- CUDA/cuDNN version: 11.7
#- GPU models and configuration: RTX 2070
#- How you installed Lightning(`conda`, `pip`, source): conda
#- Running environment of LightningApp (e.g. local, cloud): local
```

### More info

_No response_

cc @borda @justusschock @awaelchli @carmocca

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

Start with PR #16726, the migration guide, and the Trainer's fit_loop.setup_data() entry point; compare it with the former reset_train_dataloader() behavior. Run the provided minimal example on Lightning 1.9 and 2.0, then verify that changing LightningModule.train_data refreshes the loader for the next epoch and that the supported usage is documented.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.