Lightning-AI / Lightning-AI/pytorch-lightning
`trainer.fit_loop.setup_data()` does not refresh train dataset in `LightningModule`
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
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
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
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