Lightning-AI / Lightning-AI/pytorch-lightning
Resuming not correct when `max_steps` corresponds to the end of an epoch
Open
@awaelchli is already working on this.
Since Mar 12, 2024.
bug
loops
ver: 2.2.x
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
Bug description
The resuming logic for the training loop is not correct in an edge case where max_steps corresponds to the end of an epoch.
What version are you seeing the problem on?
v2.2, master
How to reproduce the bug
import shutil
import time
from pathlib import Path
import torch
from torch.utils.data import DataLoader, Dataset
from pytorch_lightning import LightningModule, Trainer
class RandomDataset(Dataset):
def __init__(self, size, length):
self.len = length
self.data = torch.randn(length, size)
def __getitem__(self, index):
return self.data[index]
def __len__(self):
return self.len
class BoringModel(LightningModule):
def __init__(self):
super().__init__()
self.layer = torch.nn.Linear(32, 2)
def forward(self, x):
return self.layer(x)
def training_step(self, dataloader_iter):
batch, _, _ = next(dataloader_iter)
batch, _, _ = next(dataloader_iter)
loss = self(batch).sum()
time.sleep(0.5)
return {"loss": loss}
def configure_optimizers(self):
return torch.optim.SGD(self.layer.parameters(), lr=0.1)
train_data = DataLoader(RandomDataset(32, 40), batch_size=2) # length = 20
checkpoint_dir = Path.cwd() / "checkpoints"
if checkpoint_dir.is_dir():
shutil.rmtree(checkpoint_dir)
model = BoringModel()
trainer = Trainer(
default_root_dir=checkpoint_dir,
max_steps=10,
# max_epochs=1,
enable_model_summary=False,
accelerator="cpu",
)
trainer.fit(model, train_data)
# Resume
trainer = Trainer(
default_root_dir=checkpoint_dir,
max_steps=20,
# max_epochs=2,
enable_model_summary=False,
accelerator="cpu",
)
trainer.fit(model, train_data, ckpt_path=checkpoint_dir / "lightning_logs/version_0/checkpoints/epoch=0-step=10.ckpt")
Error messages and logs
The expected behavior is that the loop/progress bar behaves the same when setting max_epochs (commented out above). Instead, when reaching max_steps, the epoch doesn't get bumped and when resuming the checkpoint, we resume in a mid-epoch state.
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 @carmocca @justusschock
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.