Lightning-AI / Lightning-AI/pytorch-lightning

KeyError: 'Trying to restore optimizer state but checkpoint contains only the model. This is probably due to `ModelCheckpoint.save_weights_only` being set to `True`.' But optim_cfg is in model

Open
#20,230 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug checkpointing ver: 2.2.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

I want to train from several epochs in and I can't implement correctly. I tried training by doing the following:

```
model = DeepSpeech.load_from_checkpoint(loaded_checkpoint_path)
trainer = pl.Trainer(max_epochs=max_epoch,
callbacks =[checkpoint_callback_step,checkpoint_callback_epoch],
precision = 32,accelerator = 'auto',
gradient_clip_val= 400,
accumulate_grad_batches= 1,
enable_progress_bar=False,
log_every_n_steps = 250,
logger=neptune_logger)

data_module = DeepSpeechDataModule(data_cfg,augment_cfg)
if train_from_checkpoint:
trainer.fit(model, datamodule=data_module)
```

And it treats this as epoch 0 but has the weights loaded.

If instead I do:

```
model = DeepSpeech(
labels=data_cfg.vocab, #labels,
model_cfg=model_cfg,
optim_cfg=optim_cfg,
precision=32, #DeepSpeechTrainerConf.precision,
spect_cfg=spect_cfg
)
trainer.fit(model, datamodule=data_module,ckpt_path=loaded_checkpoint_path)
```

I get the error:
KeyError: 'Trying to restore optimizer state but checkpoint contains only the model. This is probably due to `ModelCheckpoint.save_weights_only` being set to `True`.'

How do I get the saved weights and optimizer settings?

The ModelCheckpoint is:
```
checkpoint_callback_epoch = ModelCheckpoint(
dirpath=check_path,# Directory to save checkpoints
filename=f'unidirDeepSpeech_{save_targ_string}_{{epoch:02d}}-{{val_wer:.2f}}_c{{val_cer:.2f}}_{date_string}',
save_top_k=3, # Save the top 3 models
monitor='val_cer', # Metric to monitor for checkpointing
mode='min', # Minimize the monitored metric (e.g., lower val_loss is better)
save_weights_only = 'False'
)
```

Thanks for any help.

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

v2.2

### How to reproduce the bug

```python
pl.__version__ = '2.2.0.post0'
```

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current environment

```
#- PyTorch Lightning Version (e.g., 2.4.0):
#- PyTorch Version (e.g., 2.4):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
```

### 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 ModelCheckpoint configuration and the trainer.fit calls shown in the report, comparing load_from_checkpoint with ckpt_path restoration. Reproduce the KeyError on version 2.2.0.post0 and determine whether the checkpoint contains optimizer state; done means the report has a minimal reproduction and a verified restoration path.

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.