Lightning-AI / Lightning-AI/pytorch-lightning
Validation metrics not available when resuming training from checkpoint
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
I am training a lot of models in a preemtible environment where training is often stopped unexpectedly due to other jobs with higher priority requesting the GPU the model is being trained on. In these cases, I deploy resume jobs which resume from the last model checkpoint as follows:
```
trainer = Trainer(...)
datamodule = MyPLDataModule(...)
checkpoint_path = ...
module = MyPLModule.load_from_checkpoint(checkpoint_path)
trainer.fit(module, datamodule=datamodule, ckpt_path=checkpoint_path)
```
In ~90% of cases, this works perfectly fine and training resumes as expected. However, in the remaining 10%, training crashes as soon as the EarlyStoppingCallback is invoked since the monitored metric is not available.
Note that I am using a modified version of the EarlyStoppingCallback that allows for a warm-up grace-period:
```
class EarlyStoppingWithWarmUp(EarlyStopping):
def __init__(
self,
monitor: str,
min_delta: float = 0.0,
patience: int = 3,
verbose: bool = False,
mode: str = "min",
strict: bool = True,
check_finite: bool = True,
stopping_threshold: Optional[float] = None,
divergence_threshold: Optional[float] = None,
check_on_train_epoch_end: Optional[bool] = None,
log_rank_zero_only: bool = False,
warmup_epochs: int = 0,
):
super().__init__(
monitor,
min_delta,
patience,
verbose,
mode,
strict,
check_finite,
stopping_threshold,
divergence_threshold,
check_on_train_epoch_end,
log_rank_zero_only,
)
self._warmup_epochs = warmup_epochs
def _is_warmup_epoch(self, trainer: pl.Trainer) -> bool:
return trainer.current_epoch < self._warmup_epochs
def _should_skip_check_due_to_warmup(self, trainer: pl.Trainer) -> bool:
return self._is_warmup_epoch(trainer)
def on_train_epoch_end(
self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
logs = trainer.callback_metrics
if (
trainer.fast_dev_run
or not self._validate_condition_metric( # disable early_stopping with fast_dev_run
logs
)
): # short circuit if metric not present
return
if self.check_finite:
# If the monitor is not finite, run the check anyway to stop training
current = logs[self.monitor].squeeze()
if not torch.isfinite(current):
return super().on_train_epoch_end(trainer, pl_module)
if self._should_skip_check_due_to_warmup(trainer):
return
return super().on_train_epoch_end(trainer, pl_module)
def on_validation_end(
self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
if self._should_skip_check_due_to_warmup(trainer):
return
return super().on_validation_end(trainer, pl_module)
```
Note that this custom callback does not produce any issues when training a model from start to finish without resuming.
Any insights would be appreciated!
### What version are you seeing the problem on?
v2.0
### How to reproduce the bug
_No response_
### Error messages and logs
```
Early stopping conditioned on metric `val_TimeSeriesMAE` which is not available. Pass in or modify your `EarlyStopping` callback to use any of the following: `train_loss`, `train_TimeSeriesMAE`, `train_TimeSeriesMAPE`, `train_TimeSeriesRMSE`, `train_TimeSeriesMSE`
Traceback (most recent call last):
File "/app/pbim_models/cli/resume.py", line 199, in resume
trainer.fit(
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 529, in fit
call._call_and_handle_interrupt(
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 42, in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 568, in _fit_impl
self._run(model, ckpt_path=ckpt_path)
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 973, in _run
results = self._run_stage()
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1016, in _run_stage
self.fit_loop.run()
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 202, in run
self.on_advance_end()
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 369, in on_advance_end
call._call_callback_hooks(trainer, "on_train_epoch_end", monitoring_callbacks=True)
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 193, in _call_callback_hooks
fn(trainer, trainer.lightning_module, *args, **kwargs)
File "/app/models/util.py", line 110, in on_train_epoch_end
or not self._validate_condition_metric( # disable early_stopping with fast_dev_run
File "/usr/local/lib/python3.10/site-packages/pytorch_lightning/callbacks/early_stopping.py", line 146, in _validate_condition_metric
raise RuntimeError(error_msg)
RuntimeError: Early stopping conditioned on metric `val_TimeSeriesMAE` which is not available. Pass in or modify your `EarlyStopping` callback to use any of the following: `train_loss`, `train_TimeSeriesMAE`, `train_TimeSeriesMAPE`, `train_TimeSeriesRMSE`, `train_TimeSeriesMSE`
```
### Environment
Current environment
* Lightning:
- lightning-utilities: 0.9.0
- pytorch-lightning: 2.0.6
- torch: 2.0.0
- torch-cluster: 1.6.1
- torchmetrics: 0.11.4
### More info
_No response_
cc @lantiga
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 the trainer.fit resume path and the EarlyStopping._validate_condition_metric call shown in the traceback, comparing normal training with training resumed from a checkpoint. The report names no repository test or reproducible setup; the work is done when the resumed case reliably handles the monitored validation metric without the reported crash.
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
- Needs clarification
- Newbie friendliness
- 25/100