Lightning-AI/pytorch-lightning

trainer.test() with given checkpoint logs last epoch instead of checkpoint epoch

Open

#20,052 opened on Jul 5, 2024

View on GitHub
 (1 comment) (1 reaction) (0 assignees)Python (26,687 stars) (3,233 forks)batch import
bughelp wantedrepro needed

Description

Bug description

Testing from a given checkpoint leads to logging the epoch number of the last checkpoint instead of the checkpoint specified:

trainer = Trainer(..., max_epochs=10)
lightning_module = MyLightningModule(...)
datamodule = MyDatamodule()

trainer.fit(lightning_module , datamodule=datamodule)

trainer.test(lightning_module , datamodule=datamodule, ckpt_path="last")     # <-- ok: logs correct epoch and step
ckpt_path="/.../checkpoints/epoch=2-step=396.ckpt"
trainer.test(lightning_module , datamodule=datamodule, ckpt_path=ckpt_path)  # <-- incorrect: logs last epoch and step

The second test logs epoch 10 instead of epoch 2. Similarly, the step number of the second test is incorrect.

What version are you seeing the problem on?

v2.2.1

Contributor guide

Tech stack
python
Domain
backendobservability
Issue type
bug
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
3
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
1-3 hours
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
fresh
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
clear
Prerequisites
Basic knowledge of PyTorch LightningUnderstanding of checkpoint loading
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
70
Research direction
The bug is in the trainer.test() method where logging uses the global epoch/step instead of the checkpoint's epoch/step. Likely the fix is in pytorch lightning/trainer/trainer.py or the logging callback (pytorch lightning/callbacks/model checkpoint.py). Check how ckpt path is used to restore state and ensure that the logger's current epoch and global step are set from the checkpoint. The issue has clear reproduction steps and no assignees, making it a good candidate for a first time contributor.