Step-by-step: Logging custom metrics from a PyTorch Lightning callback.

devops_dad_joke · 2026-07-29T01:35:56Z

Alright, gather 'round the campfire, kids. Let's talk about making your PyTorch Lightning trainer actually tell W&B what it's doing, beyond the default auto-logged stuff. Because let's be honest, the auto-logging is great until you need to track something weird, like the gradient norm of a specific layer or a custom performance metric on a weird validation subset. Then you're left scratching your head, digging through callbacks. Here's the deal: you need a custom callback. But the trick isn't just *having* the callback—it's making sure you log to the right experiment run inside it. If you just call `wandb.log()` directly from a Lightning callback, you might be yelling into the void (or worse, logging to a different run). Here's my go-to pattern. I create a `WandbMetricsLogger` callback that grabs the current W&B run from the logger. PyTorch Lightning 2.x+ makes this pretty clean. ```python import pytorch_lightning as pl import torch import wandb class WandbCustomMetricsCallback(pl.Callback): def on_validation_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None: # This is the key part: get the wandb logger instance wandb_logger = None for logger in trainer.loggers: if isinstance(logger, pl.loggers.WandbLogger): wandb_logger = logger break if wandb_logger is not None: # Calculate your custom metric. Example: mean absolute value of some layer's weights custom_metric = pl_module.some_layer.weight.abs().mean().item() # Log using the logger's experiment, which is the wandb.Run object wandb_logger.experiment.log({ "custom/weight_norm": custom_metric, "epoch": trainer.current_epoch }) ``` Then you just slot it into your trainer: ```python trainer = pl.Trainer( callbacks=[WandbCustomMetricsCallback()], logger=pl.loggers.WandbLogger(project="my_project") ) ``` Why do it this way? Because the `WandbLogger` manages the run lifecycle. If you instantiate a `wandb.init()` inside your callback, you're asking for a fight with duplicate runs or dropped logs. Let the logger handle the plumbing, you just hand it the data. The pitfall I see all the time is folks trying to import and use `wandb` globally in the callback without ensuring they're in the right run context. This pattern ties your logging directly to the run the Lightning trainer started. No more lost metrics. Now go forth and log that weird, bespoke metric your PM insists on tracking. Just don't blame me when your W&B bill goes up. - tm

Page 2 / 2 Prev

Weights and Biases Reviews

Last Post by isabelc 1 hour ago

16 Posts

16 Users

0 Reactions

4 Views

RSS

isabelc

(@isabelc)

Eminent Member

Joined: 1 week ago

Posts: 18

29/07/2026 2:50 pm

I see what you mean about building a trap. That silent failure when someone adds a second logger is exactly the kind of thing that would take me days to debug later. So the real first step is a more defensive check to find the right logger.

Do you have a go-to snippet for that check? Something that reliably handles the LoggerCollection case?

ReplyQuote

Page 2 / 2 Prev