Lightning-AI / Lightning-AI/pytorch-lightning

Restoring a checkpoint (from model checkpoint) via `trainer.test(model_module, ckpt_path='best')` doesn't restore associated `current_epoch` within the trainer

Open
#19,558 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug checkpointing trainer: test
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

When a checkpoint (presumably the 'best' one saved by model checkpoint monitor) is restored to be used for testing via `trainer.test(model_module, ckpt_path='best')` the `current_epoch` member of the trainer is not restored to the one associated with the checkpoint, even though the checkpoint file stores the correct epoch value.

I am not sure whether this is intended (seems rather incorrect to me). I asked a question about this problem at the Discord forum before, but got no response, so I decided to open an issue here.

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

v2.2

### How to reproduce the bug

```python
import os

import torch
from lightning.pytorch import LightningModule, Trainer, callbacks, seed_everything
from torch.utils.data import DataLoader, Dataset

class RandomDataset(Dataset):
def __init__(self, size, length, offset = 0):
self.len = length
self.data = torch.randn(length, size) + offset

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, batch, batch_idx):
loss = self(batch).sum()
self.log("train_loss", loss)
return {"loss": loss}

def validation_step(self, batch, batch_idx):
loss = self(batch).sum()
self.log("valid_loss", loss)

def test_step(self, batch, batch_idx):
loss = self(batch).sum()
print(f'CURRENT EPOCH: {self.current_epoch}')
self.log("test_loss", loss)

def configure_optimizers(self):
return torch.optim.SGD(self.layer.parameters(), lr=0.1)

def run():
seed_everything(10)

train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
val_data = DataLoader(RandomDataset(32, 64, offset=16), batch_size=2)
test_data = DataLoader(RandomDataset(32, 64), batch_size=2)

model = BoringModel()
trainer = Trainer(
default_root_dir=os.getcwd(),
num_sanity_val_steps=0,
limit_train_batches=1,
limit_val_batches=1,
limit_test_batches=1,
max_epochs=2,
enable_model_summary=False,
callbacks=[(mc := callbacks.ModelCheckpoint(monitor="valid_loss", verbose=True))],
)
trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)
trainer.test(model, dataloaders=test_data, ckpt_path="best")
print(f'CKPT EPOCH: {torch.load(mc.best_model_path, map_location="cpu")["epoch"]}')

if __name__ == "__main__":
run()
```

### Error messages and logs

```
Epoch 0: 100%|███████████████████████████████████████████| 32/32 [00:00<00:00, 277.12it/s, v_num=7]
Epoch 0, global step 32: 'valid_loss' reached 183.69585 (best 183.69585), saving model to '/Users/maciej/Desktop/lightning_logs/version_7/checkpoints/epoch=0-step=32.ckpt' as top 1
Epoch 1: 100%|███████████████████████████████████████████| 32/32 [00:00<00:00, 432.74it/s, v_num=7]
Epoch 1, global step 64: 'valid_loss' was not in top 1
`Trainer.fit` stopped: `max_epochs=2` reached.
Epoch 1: 100%|███████████████████████████████████████████| 32/32 [00:00<00:00, 410.05it/s, v_num=7]
Restoring states from the checkpoint path at /Users/maciej/Desktop/lightning_logs/version_7/checkpoints/epoch=0-step=32.ckpt
Loaded model weights from the checkpoint at /Users/maciej/Desktop/lightning_logs/version_7/checkpoints/epoch=0-step=32.ckpt
/Users/maciej/.local/share/pyenv/versions/3.11.7/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:441: The 'test_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=9` in the `DataLoader` to improve performance.
Testing DataLoader 0: 0%| | 0/1 [00:00
Current environment

```
* CUDA:
- GPU: None
- available: False
- version: None
* Lightning:
- lightning: 2.2.0.post0
- lightning-utilities: 0.10.1
- pytorch-lightning: 2.2.0.post0
- torch: 2.2.1
- torchmetrics: 1.3.1
- torchvision: 0.17.1
* Packages:
- aiohttp: 3.9.3
- aiohttp-retry: 2.8.3
- aiosignal: 1.3.1
- amqp: 5.2.0
- annotated-types: 0.6.0
- antlr4-python3-runtime: 4.9.3
- anyio: 4.3.0
- appdirs: 1.4.4
- appnope: 0.1.4
- argon2-cffi: 23.1.0
- argon2-cffi-bindings: 21.2.0
- arrow: 1.3.0
- astroid: 3.1.0
- asttokens: 2.4.1
- async-lru: 2.0.4
- asyncssh: 2.14.2
- atpublic: 4.0
- attrs: 23.2.0
- autopep8: 2.0.4
- babel: 2.14.0
- backcall: 0.2.0
- beautifulsoup4: 4.12.3
- billiard: 4.2.0
- black: 24.2.0
- bleach: 6.1.0
- blinker: 1.7.0
- celery: 5.3.6
- certifi: 2024.2.2
- cffi: 1.16.0
- charset-normalizer: 3.3.2
- click: 8.1.7
- click-didyoumean: 0.3.0
- click-plugins: 1.1.1
- click-repl: 0.3.0
- colorama: 0.4.6
- comm: 0.2.1
- configobj: 5.0.8
- contourpy: 1.2.0
- cryptography: 42.0.5
- cycler: 0.12.1
- dash: 2.15.0
- dash-core-components: 2.0.0
- dash-html-components: 2.0.0
- dash-table: 5.0.0
- debugpy: 1.8.1
- decorator: 5.1.1
- defusedxml: 0.7.1
- dictdiffer: 0.9.0
- dill: 0.3.8
- diskcache: 5.6.3
- distlib: 0.3.8
- distro: 1.9.0
- docopt: 0.6.2
- dpath: 2.1.6
- dulwich: 0.21.7
- dvc: 3.48.0
- dvc-data: 3.13.0
- dvc-http: 2.32.0
- dvc-objects: 5.0.0
- dvc-render: 1.0.1
- dvc-studio-client: 0.20.0
- dvc-task: 0.3.0
- einops: 0.7.0
- entrypoints: 0.4
- executing: 2.0.1
- fastjsonschema: 2.19.1
- filelock: 3.13.1
- flake8: 7.0.0
- flask: 3.0.2
- flatten-dict: 0.4.2
- flufl.lock: 7.1.1
- fonttools: 4.49.0
- fqdn: 1.5.1
- frozenlist: 1.4.1
- fsspec: 2024.2.0
- funcy: 2.0
- gitdb: 4.0.11
- gitpython: 3.1.42
- grandalf: 0.8
- greenlet: 3.0.3
- gto: 1.7.0
- h11: 0.14.0
- h5py: 3.10.0
- httpcore: 1.0.4
- httpx: 0.27.0
- hydra-core: 1.3.2
- idna: 3.6
- imageio: 2.34.0
- importlib-metadata: 7.0.1
- iniconfig: 2.0.0
- ipdb: 0.13.13
- ipykernel: 6.29.3
- ipython: 8.12.3
- isoduration: 20.11.0
- isort: 5.13.2
- iterative-telemetry: 0.0.8
- itsdangerous: 2.1.2
- jedi: 0.19.1
- jinja2: 3.1.3
- joblib: 1.3.2
- json5: 0.9.17
- jsonpointer: 2.4
- jsonschema: 4.21.1
- jsonschema-specifications: 2023.12.1
- jupyter-client: 8.6.0
- jupyter-core: 5.7.1
- jupyter-events: 0.9.0
- jupyter-lsp: 2.2.3
- jupyter-server: 2.12.5
- jupyter-server-mathjax: 0.2.6
- jupyter-server-terminals: 0.5.2
- jupyterlab: 4.1.2
- jupyterlab-pygments: 0.3.0
- jupyterlab-server: 2.25.3
- kiwisolver: 1.4.5
- kombu: 5.3.5
- lazy-loader: 0.3
- lightning: 2.2.0.post0
- lightning-utilities: 0.10.1
- markdown-it-py: 3.0.0
- markupsafe: 2.1.5
- matplotlib: 3.8.3
- matplotlib-inline: 0.1.6
- mccabe: 0.7.0
- mdurl: 0.1.2
- mistune: 3.0.2
- mpmath: 1.3.0
- msgpack: 1.0.7
- multidict: 6.0.5
- mypy: 1.8.0
- mypy-extensions: 1.0.0
- nbclient: 0.9.0
- nbconvert: 7.16.1
- nbdime: 4.0.1
- nbformat: 5.9.2
- neovim: 0.3.1
- nest-asyncio: 1.6.0
- networkx: 3.2.1
- notebook-shim: 0.2.4
- numpy: 1.26.4
- omegaconf: 2.3.0
- orjson: 3.9.15
- overrides: 7.7.0
- packaging: 23.2
- pandas: 2.2.1
- pandocfilters: 1.5.1
- parso: 0.8.3
- pathspec: 0.12.1
- pexpect: 4.9.0
- pickleshare: 0.7.5
- pillow: 10.2.0
- pip: 23.2.1
- pipreqs: 0.5.0
- platformdirs: 3.11.0
- plotly: 5.19.0
- pluggy: 1.4.0
- prometheus-client: 0.20.0
- prompt-toolkit: 3.0.43
- psutil: 5.9.8
- ptyprocess: 0.7.0
- pure-eval: 0.2.2
- pycodestyle: 2.11.1
- pycparser: 2.21
- pydantic: 2.6.3
- pydantic-core: 2.16.3
- pydocstyle: 6.3.0
- pydot: 2.0.0
- pyflakes: 3.2.0
- pygit2: 1.14.1
- pygments: 2.17.2
- pygtrie: 2.5.0
- pylint: 3.1.0
- pynvim: 0.5.0
- pyparsing: 3.1.1
- pytest: 8.0.2
- pytest-mock: 3.12.0
- python-dateutil: 2.8.2
- python-json-logger: 2.0.7
- pytoolconfig: 1.3.1
- pytorch-lightning: 2.2.0.post0
- pytz: 2024.1
- pyyaml: 6.0.1
- pyzmq: 25.1.2
- referencing: 0.33.0
- rentry: 1.0.1
- requests: 2.31.0
- retrying: 1.3.4
- rfc3339-validator: 0.1.4
- rfc3986-validator: 0.1.1
- rich: 13.7.0
- rope: 1.12.0
- rpds-py: 0.18.0
- ruamel.yaml: 0.18.6
- ruamel.yaml.clib: 0.2.8
- scikit-image: 0.22.0
- scikit-learn: 1.4.1.post1
- scipy: 1.12.0
- scmrepo: 3.1.0
- semver: 3.0.2
- send2trash: 1.8.2
- setuptools: 65.5.0
- shortuuid: 1.0.11
- shtab: 1.7.0
- six: 1.16.0
- smmap: 5.0.1
- sniffio: 1.3.1
- snowballstemmer: 2.2.0
- soupsieve: 2.5
- sqltrie: 0.11.0
- stack-data: 0.6.3
- sympy: 1.12
- tabulate: 0.9.0
- tenacity: 8.2.3
- terminado: 0.18.0
- threadpoolctl: 3.3.0
- tifffile: 2024.2.12
- tinycss2: 1.2.1
- tomlkit: 0.12.4
- torch: 2.2.1
- torchmetrics: 1.3.1
- torchvision: 0.17.1
- tornado: 6.4
- tqdm: 4.66.2
- traitlets: 5.14.1
- typer: 0.9.0
- types-python-dateutil: 2.8.19.20240106
- typing-extensions: 4.10.0
- tzdata: 2024.1
- uri-template: 1.3.0
- urllib3: 2.2.1
- vine: 5.1.0
- virtualenv: 20.25.1
- voluptuous: 0.14.2
- wcwidth: 0.2.13
- webcolors: 1.13
- webencodings: 0.5.1
- websocket-client: 1.7.0
- werkzeug: 3.0.1
- wheel: 0.42.0
- yarg: 0.1.9
- yarl: 1.9.4
- zc.lockfile: 3.0.post1
- zipp: 3.17.0
* System:
- OS: Darwin
- architecture:
- 64bit
-
- processor: arm
- python: 3.11.7
- release: 23.2.0
- version: Darwin Kernel Version 23.2.0: Wed Nov 15 21:53:18 PST 2023; root:xnu-10002.61.3~2/RELEASE_ARM64_T6000
```

### More info

_No response_

cc @lantiga

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 the Trainer.test call using ckpt_path='best', the ModelCheckpoint callback, and the test_step that prints current_epoch. Trace how the selected checkpoint's stored epoch is loaded during testing, then add a regression test using the supplied reproduction to verify that current_epoch matches the checkpoint value.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.