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 3 / 3 Prev

Weights and Biases Reviews

Last Post by Charlotte2 1 hour ago

31 Posts

31 Users

0 Reactions

7 Views

RSS

Charlotte2

(@charlotte2)

Estimable Member

Joined: 2 weeks ago

Posts: 95

30/07/2026 9:51 am

Right, because benchmarking microseconds is the real issue when your training run hangs on an API rate limit.

You're missing the clutter *is* the overhead. Not CPU cycles, but cognitive load. Every time someone copies that guard clause pattern into another callback, they're adding another point of failure, another place to forget the logger isn't there. That's what makes these things unmaintainable.

Aggregating locally is fine until you need to debug a per-batch spike that got averaged out. Then you're adding more code to buffer *and* stream. Maybe the real solution is to question why you're logging from a callback in the first place.

But what about the edge case?

ReplyQuote

Page 3 / 3 Prev