Lightning-AI / Lightning-AI/pytorch-lightning

Periodic checkpoints logged for every epoch to wandb with ModelCheckpoint & WandbLogger

Open
#17,916 0 comments 2 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug callback: model checkpoint logger: wandb ver: 2.0.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

Hello,
I am using `ModelCheckpoint` along with `WandbLogger` and I want to create snapshot checkpoints, i.e. create a checkpoint every `N` epochs. When I create my checkpoint callback and logger as given in the below snippet, the checkpoints created in my local machine are correct, i.e. there are checkpoints corresponding to epoch 49, 99, and so on. However, in my wandb project there is an artifact entry for each epoch and as a consequence both local wandb cache directory and wandb storage usage grows quickly.

I believe that the documentation of `ModelCheckpoint` is unclear about how to achieve periodic checkpoints because providing only `every_n_epochs=50` alone or with `save_top_k=3, monitor="epoch"` only created one checkpoint in my case. Is there any other way to achieve periodic checkpointing?

I am experiencing this issue because of the `save_last=True` argument combined with the implementation of the `WandbLogger` which calls the following code block after checkpoints. I believe the direct upload of the last checkpoint after each epoch contradicts the described behavior in the documentation, and these checkpoints should only be uploaded after training is completed with `log_model=True`. There seems to be no way to get periodic checkpoints along with the last checkpoint before training ends with the current behavior.
```python
# log checkpoints as artifacts
if self._log_model == "all" or self._log_model is True and checkpoint_callback.save_top_k == -1:
self._scan_and_log_checkpoints(checkpoint_callback)
```

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

v2.0

### How to reproduce the bug

```python
import pytorch_lightning as L

wandb_logger = L.loggers.WandbLogger(
entity=ENTITY,
save_dir=LOG_DIR,
log_model=True, # uploads the checkpoints that are created at the end of the training
)

checkpoint_callback = L.callbacks.ModelCheckpoint(
dirpath=Path(LOG_DIR).joinpath(ID, "checkpoints"),
save_last=True,
save_top_k=-1,
every_n_epochs=50
)
```

### Error messages and logs

_No response_

### Environment

Current environment

```
#- Lightning Component: Trainer, LightningModule
#- PyTorch Lightning Version: 2.0.3
#- PyTorch Version (e.g., 2.0): 2.0.1
#- Python version (e.g., 3.9): 3.11.2
#- Wandb version: 0.15.4
#- OS (e.g., Linux): Linux
#- How you installed Lightning(`conda`, `pip`, source): conda/micromamba
#- Running environment of LightningApp (e.g. local, cloud): local
```

### More info

_No response_

cc @lantiga @morganmcg1 @borisdayma @scottire @parambharat

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 and WandbLogger entry points in the v2.0 implementation, especially the _scan_and_log_checkpoints call shown in the issue. Reproduce with every_n_epochs=50, save_last=True, save_top_k=-1, and log_model=True; done means periodic local checkpoints do not create an artifact for every epoch, with the resulting behavior documented and tested.

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
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.