WandbLogger(log_model=False) does not work, model is always saved to wandb logs

[N0vel]

🐛 Bug

WandbLogger log_model argument does not work as expected. The model checkpoint is always being saved to wandb logs (project_name/run/checkpoints/checkpoint_name.ckpt). The model is not uploaded to wandb, only saved on the drive. It should not be saved on disk nor uploaded.
It is especially important during sweeping (hundreds of runs exhaust available space quickly). I don't use ModelCheckpoint here so the weights should not be saved at all.

To Reproduce

wandb_logger = WandbLogger(project='private_example', config={
**model_args.dict, **data_args.dict, **training_args.dict
}, log_model=False)
trainer = pl.Trainer(...,
logger=[wandb_logger])

Expected behavior

Checkpoints are saved if log_model True and not saved otherwise.

Environment

• PyTorch Version (e.g., 1.0): 1.7.1
• OS (e.g., Linux): Debian
• How you installed PyTorch (conda, pip, source): conda
• Python version: 3.6.9
• CUDA/cuDNN version: cuda 11
• wandb version: 0.10.15 (same problem on 0.10.17)
• lightning version: 1.1.4 (same problem on 1.1.8)

[awaelchli]

We use wandb.save here, so it is normal that the file ends up on the disk first. It should get uploaded only at the end of the run. Maybe your process gets killed and wandb with it too, it would explain why artifacts are not uploaded

[N0vel]

We use wandb.save here, so it is normal that the file ends up on the disk first. It should get uploaded only at the end of the run. Maybe your process gets killed and wandb with it too, it would explain why artifacts are not uploaded

It should not be saved on disk nor uploaded. Is it possible with WandbLogger? I just don't want 500 checkpoints on SSD after wandb sweep...

[N0vel]

Lightning documentation says 'log_model (bool) – Save checkpoints in wandb dir to upload on W&B servers.'
So if it's False, why should checkpoints be saved in wandb dir?

[awaelchli]

sorry, you are right and I misread your report first thinking it doesn't upload.

[awaelchli]

1. Checkpointing is turned on by default. You can turn it off with checkpoint_callbacks=False.

2. It looks like wandb will upload everything that is in the default logging directory under wandb.
   If you set save_dir="./somewhere" you can prevent this.

[N0vel]

Thank you, looks like checkpoint_callbacks=False works right (though WandbLogger documentation is misleading).

[dlangerm]

This is still a bug and should not be closed. The documentation and semantics for the init are very misleading. If log_model is false then the model should not be uploaded.

[awaelchli]

@dlangerm my answer still applies for the problem that @N0vel had.
You are reporting that log_model=False will still upload files to wandb right?
Then let's investigate that?

[dlangerm]

@awaelchli Yes log_model=False still causes the model to be uploaded at the end of the run, IMO if the desired behavior is clearly not to log models to WandB the flag should be passed to the wandb module to prevent it from uploading the model. Otherwise, there is no obvious way to prevent the model from being uploaded through the torch lightning interface other than setting save_dir="./somewhere" which appears to only work on some platforms (this seems to work on Linux but on Windows, the model is still uploaded but I haven't tested it all that thoroughly).

[awaelchli]

I tried this:

import os
from pathlib import Path

import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer
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("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=2)

    model = BoringModel()
    logger = WandbLogger(project="logged-model", log_model=False)
    trainer = Trainer(
        logger=logger,
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        num_sanity_val_steps=0,
        max_epochs=1,
        weights_summary=None,
    )
    trainer.fit(model, train_dataloader=train_data)

    model_path = Path(trainer.checkpoint_callback.best_model_path)
    assert model_path.exists()
    print(os.listdir(model_path.parent))


if __name__ == '__main__':
    run()

It does not uploading any checkpoint files, at least I can't see them in the files tab on the experiment.
Can you try to adapt this example please to help me reproduce what you are seeing?

wandb version: 0.10.32
PL version: master
macos