Lightning-AI / Lightning-AI/pytorch-lightning
WandbLogger does not log checkpoints from multiple ModelCheckpoint callbacks, but only from last one
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
## 🐛 Bug
As the title explains, WandbLogger does not log checkpoints from multiple ModelCheckpoint callbacks, but only from last one. When multiple ModelCheckpoint callbacks are used, only the last one provided to the trainer in `callbacks=[..., LastModelCheckpoint]` is logged in Wandb, even though all the checkpoints are created as expected locally. Is this intended to be like this or is it an issue?
### To Reproduce
```
import os
import torch
from torch.utils.data import DataLoader, Dataset
from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.loggers import WandbLogger
class RandomDataset(Dataset):
def __init__(self, size, length):
self.len = length
self.data = torch.randn(length, size)
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("val_loss", loss)
def test_step(self, batch, batch_idx):
loss = self(batch).sum()
self.log("test_loss", loss)
def configure_optimizers(self):
return torch.optim.SGD(self.layer.parameters(), lr=0.1)
def run():
train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
val_data = DataLoader(RandomDataset(32, 64), batch_size=2)
test_data = DataLoader(RandomDataset(32, 64), batch_size=2)
model = BoringModel()
max_steps = 500
min_epochs = 10
callback_last_step = ModelCheckpoint(
verbose=False,
save_on_train_epoch_end=False,
dirpath="./checkpoints",
filename="{epoch}_{step}_{val_loss:.2f}_last_step",
every_n_train_steps=max_steps,
)
callback_epoch10 = ModelCheckpoint(
verbose=False,
save_on_train_epoch_end=False,
dirpath="./checkpoints",
filename="{epoch}_{step}_{val_loss:.2f}_epoch10",
every_n_epochs=min_epochs,
)
callback_topk = ModelCheckpoint(
verbose=False,
save_on_train_epoch_end=False,
dirpath="./checkpoints",
filename="{epoch}_{step}_{val_loss:.2f}_topk",
every_n_epochs=1,
monitor="val_loss",
mode="min",
save_top_k=3,
)
callbacks = [callback_topk, callback_last_step, callback_epoch10]
logger = WandbLogger(
project="MNIST", entity="aleqrc", name="test", log_model=True, mode="online"
)
trainer = Trainer(
min_epochs=min_epochs,
max_steps=max_steps,
callbacks=callbacks,
logger=logger,
limit_train_batches=50,
limit_val_batches=50,
enable_model_summary=False,
num_sanity_val_steps=0,
)
trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)
if __name__ == "__main__":
run()
```
### Expected behavior
The code should end successfully, create 5 checkpoints (3 ending with "topk", 1 ending with "epoch10" and 1 ending with "last_step"), and log all of them in Wandb, however only checkpoints created by the last ModelCheckpoint provided to the trainer are logged.
### Environment
- PyTorch Lightning Version: 1.5.4
- PyTorch Version: 1.10.0
- Python version: 3.8.10
- OS: Windows 11
- CUDA/cuDNN version: 11.3.1
### Additional context
I have also encountered this error on a Linux machine (Ubuntu), even though I did not test exactly this code snippet, but I have the same issue there as well.
cc @awaelchli @morganmcg1 @AyushExel @borisdayma @scottire
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.
Assessment
This issue has not been assessed yet.