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
Getting the wandb logger instance is indeed the key. But what's the actual ROI of doing this in every callback? Could get messy if you have multiple custom metrics.
I usually attach a reference to the logger in the callback's `setup` method. That way you're not fishing for it on every epoch end. Less overhead, same result.
Also, watch out for sync issues. If your validation step is fast, logging too many custom metrics can throttle W&B.
Ask me about hidden egress costs.
Yeah, less overhead. Because the microseconds you save avoiding a logger lookup each epoch are what's really bottlenecking your training loop.
The sync issue is real though. Watched a team burn a week "optimizing" their model when the real problem was W&B throttling from their overzealous callback logging every batch. The logs fell behind and skewed their perceived step times. Classic.
If you're that worried about overhead, maybe don't log custom metrics at all. The default ones probably tell you what you need to know.
If it ain't broke, don't 'upgrade' it.
Your code snippet cut off, but the critical step is indeed accessing the logger instance. The pattern you're describing works, but I'd add that you should also check if a W&B logger is actually present. It's easy to forget and then your callback crashes during a dry-run or when testing with a different logger like TensorBoard.
I usually add a guard clause at the start of the logging method: `if not isinstance(trainer.logger, WandbLogger): return`. This makes the callback more portable and avoids silent failures when the expected logger isn't configured.
Support is a product, not a department.
Your point about checking the logger type is crucial for production code. That guard clause prevents a whole class of runtime errors when someone swaps out loggers for local debugging or uses multi-logger configurations.
I'd extend it slightly by suggesting you also validate the logger's experiment attribute. Sometimes the logger exists but hasn't been initialized, particularly if you're attaching callbacks before the trainer starts. A more defensive pattern is to check `if hasattr(trainer.logger, 'experiment') and trainer.logger.experiment is not None`.
This becomes especially relevant in distributed training setups where logger initialization has subtle timing differences.