Support `wandb_logger.watch()` when using `LightningCLI`

[Boltzmachine]

Description & Motivation

When I am using LigthtningCLI, I find nowhere to set wandb_logger.watch(model)

Pitch

Perhaps adding a callback in LightningCLI will work?

Alternatives

No response

Additional context

No response

cc @lantiga @Borda @morganmcg1 @borisdayma @scottire @parambharat @mauvilsa

[carlthome]

Same issue here.

[carlthome]

class WatchModel(Callback):
    def on_fit_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
        if isinstance(trainer.logger, WandbLogger):
            trainer.logger.watch(pl_module, log="all", log_freq=10)

I'll just do this for now and it seems to work.

[mauvilsa]

class WatchModel(Callback):
def on_fit_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
if isinstance(trainer.logger, WandbLogger):
trainer.logger.watch(pl_module, log="all", log_freq=10)

I'll just do this for now and it seems to work.

@carlthome can you explain why you say "for now". Isn't this a good solution?

[carlthome]

class WatchModel(Callback):
def on_fit_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
if isinstance(trainer.logger, WandbLogger):
trainer.logger.watch(pl_module, log="all", log_freq=10)

I'll just do this for now and it seems to work.

@carlthome can you explain why you say "for now". Isn't this a good solution?

Not keen on needing custom callbacks to get this working. Unintended side effects that go untested for in the framework tend to creep in over time, and puts more maintenance burden on my colleagues once we update dependencies.

[mauvilsa]

Fair enough. I am not a user of wandb, so I don't think I could help here. But if any user of wandb is able to come up with something better, and opens a pull request including proper tests, then I guess the lightning maintainers could merge it. Or maybe better to first describe the proposed solution here. I can help refine the idea so that it fits the generic spirit of LightningCLI.

[SkafteNicki]

Since the watch method is unique to the WandbLogger:
https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.loggers.wandb.html#module-lightning.pytorch.loggers.wandb
I do not really see that it make sense to build into lightning cli. We could consider adding some kind of setup method but I do wonder if it make sense because there is such much customization that can be done by the different loggers.

In the end I do think a custom callback here is the best option, because it fits the specific need.

[SkafteNicki]

Looked further into this and my conclusion still stands that enabling in a more streamlined way is not really possible. A custom callback is still the better option here. For completeness, I will post a full script + config for future reference, that just adds the extra arguments that watch takes + makes sure watch is called with the same frequency as all other logging.

from lightning.pytorch.demos import BoringDataModule, BoringModel
from lightning.pytorch.cli import LightningCLI
from lightning.pytorch.loggers import WandbLogger
from lightning.pytorch.callbacks import Callback
from typing import Optional

class WatchModel(Callback):
    def __init__(self, log_type: Optional[str] = "gradients", log_graph: bool = True):
        self.log_type = log_type
        self.log_graph = log_graph

    def on_fit_start(self, trainer, pl_module):
        logger = trainer.logger
        if isinstance(logger, WandbLogger):
            logger.watch(pl_module, log=self.log_type, log_freq=trainer.log_every_n_steps, log_graph=self.log_graph)


if __name__ == "__main__":
    cli = LightningCLI(
        BoringModel,
        BoringDataModule,
        save_config_callback=None,
    )

trainer:
  callbacks:
    - class_path: __main__.WatchModel
      init_args:
        log_type: "gradients"
        log_graph: true
  logger:
    - class_path: lightning.pytorch.loggers.WandbLogger
      init_args:
        project: "test-project"
  log_every_n_steps: 50
  max_epochs: 3

Closing issue.