PyTorch Lightning

PyTorch Lightning provides a lightweight wrapper for organizing your PyTorch code and easily adding advanced features such as distributed training and 16-bit precision. W&B provides a lightweight wrapper for logging your ML experiments. We're incorporated directly into the PyTorch Lightning library, so you can always check out their documentation.

⚡Get going lightning-fast with just two lines:

from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning import Trainer
wandb_logger = WandbLogger()
trainer = Trainer(logger=wandb_logger)

✅ Check out real examples!

We've created a few examples for you to see how the integration works:

💻 API Reference

WandbLogger

Parameters:

  • name (str) – display name for the run.

  • save_dir (str) – path where data is saved.

  • offline (bool) – run offline (data can be streamed later to wandb servers).

  • version (id) – sets the version, mainly used to resume a previous run.

  • anonymous (bool) – enables or explicitly disables anonymous logging.

  • project (str) – the name of the project to which this run will belong.

  • tags (list of str) – tags associated with this run.

WandbLogger.watch

Log model topology as well as optionally gradients and weights.

wandb_logger.watch(model, log='gradients', log_freq=100)

Parameters:

  • model (nn.Module) – model to be logged.

  • log (str) – can be "gradients" (default), "parameters", "all" or None.

  • log_freq (int) – step count between logging of gradients and parameters.

WandbLogger.log_hyperparams

Record hyperparameter configuration.

Note: this function is called automatically by Trainer

wandb_logger.log_hyperparams(params)

Parameters:

  • params (dict) – dictionary with hyperparameter names as keys and configuration values as values

WandbLogger.log_metrics

Record training metrics.

Note: this function is called automatically by Trainer

wandb_logger.log_metrics(metrics, step=None)

Parameters:

  • metric (numeric) – dictionary with metric names as keys and measured quantities as values

  • step (int|None) – step number at which the metrics should be recorded