Lightning-AI / Lightning-AI/pytorch-lightning

WandbLogger does not log checkpoints from multiple ModelCheckpoint callbacks, but only from last one

Open
#11,156 6 comments 0 reactions 1 assignee View on GitHub

Nobody has claimed this yet.

bug logger: wandb
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

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.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.