Wandb-dependent Model Checkpoint

[kelvins64]

I'm trying to set my model checkpoint path based on the Wandb experiment while using DDP. However, we only have access to the actual Wandb experiment in rank 0 (see this discussion). Therefore, trying to base the path on logger.experiment will fail as it will be a DummyExperiment:

import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.loggers import WandbLogger

logger = WandbLogger()
ckpt_callback = ModelCheckpoint(dirpath=do_something(logger.experiment)) # can't operate on actual logger.experiment
kwargs = {
    'logger': logger,
    'gpus': 2,
    'strategy': 'ddp',
    'callbacks': ckpt_callback
}
trainer = pl.Trainer(**kwargs)

On the other hand, DDP hangs if I try to detect whether the process is rank zero, and add the ModelCheckpoint to callbacks (or simply change its dirpath) based on it:

import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.loggers.base import DummyExperiment

def is_rank_zero(wandb_logger: WandbLogger):
    return not isinstance(wandb_logger.experiment, DummyExperiment)

logger = WandbLogger()
kwargs = {
    'logger': logger,
    'gpus': 2,
    'strategy': 'ddp'
}

if is_rank_zero(logger):
    ckpt_callback = ModelCheckpoint(dirpath=do_something(logger.experiment))
    kwargs['callbacks'] = ckpt_callback

trainer = pl.Trainer(**kwargs)

What is the correct way to build the ModelCheckpoint directory path based on the Wandb experiment?

Full example with BoringModel

import torch
from torch.utils.data import DataLoader, Dataset
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.loggers.base import DummyExperiment

def is_rank_zero(wandb_logger: WandbLogger):
    return not isinstance(wandb_logger.experiment, DummyExperiment)

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(pl.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("valid_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=16)
    val_data = DataLoader(RandomDataset(32, 64), batch_size=16)

    model = BoringModel()

    logger = WandbLogger()
    kwargs = {
        'logger': logger,
        'gpus': 2,
        'strategy': 'ddp',
        'max_epochs': 1
    }

    if is_rank_zero(logger):
        dirpath = logger.experiment.dir.split('wandb/')[1].split('/')[0] # run-XXXXX
        ckpt_callback = ModelCheckpoint(dirpath=dirpath)
        kwargs['callbacks'] = ckpt_callback

    trainer = pl.Trainer(**kwargs)
    trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)

if __name__ == "__main__":
    run()