lightning_template.utils.loggers.wandb

lightning_template.utils.loggers.wandb#

Classes#

WandbNamedLogger

Module Contents#

class lightning_template.utils.loggers.wandb.WandbNamedLogger(name: str | None = None, save_dir: lightning.fabric.utilities.types._PATH = '.', version: str | None = None, offline: bool = False, dir: lightning.fabric.utilities.types._PATH | None = None, id: str | None = None, anonymous: bool | None = None, project: str | None = None, log_model: Literal['all'] | bool = False, experiment: wandb.wandb_run.Run | wandb.sdk.lib.RunDisabled | None = None, prefix: str = '', checkpoint_name: str | None = None, **kwargs: Any)#

Bases: lightning.pytorch.loggers.wandb.WandbLogger

Log using Weights and Biases.

Installation and set-up

Install with pip:

pip install wandb

Create a WandbLogger instance:

from lightning.pytorch.loggers import WandbLogger

wandb_logger = WandbLogger(project="MNIST")

Pass the logger instance to the Trainer:

trainer = Trainer(logger=wandb_logger)

A new W&B run will be created when training starts if you have not created one manually before with wandb.init().

Log metrics

Log from LightningModule:

class LitModule(LightningModule):
    def training_step(self, batch, batch_idx):
        self.log("train/loss", loss)

Use directly wandb module:

wandb.log({"train/loss": loss})

Log hyper-parameters

Save LightningModule parameters:

class LitModule(LightningModule):
    def __init__(self, *args, **kwarg):
        self.save_hyperparameters()

Add other config parameters:

# add one parameter
wandb_logger.experiment.config["key"] = value

# add multiple parameters
wandb_logger.experiment.config.update({key1: val1, key2: val2})

# use directly wandb module
wandb.config["key"] = value
wandb.config.update()

Log gradients, parameters and model topology

Call the watch method for automatically tracking gradients:

# log gradients and model topology
wandb_logger.watch(model)

# log gradients, parameter histogram and model topology
wandb_logger.watch(model, log="all")

# change log frequency of gradients and parameters (100 steps by default)
wandb_logger.watch(model, log_freq=500)

# do not log graph (in case of errors)
wandb_logger.watch(model, log_graph=False)

The watch method adds hooks to the model which can be removed at the end of training:

wandb_logger.experiment.unwatch(model)

Log model checkpoints

Log model checkpoints at the end of training:

wandb_logger = WandbLogger(log_model=True)

Log model checkpoints as they get created during training:

wandb_logger = WandbLogger(log_model="all")

Custom checkpointing can be set up through ModelCheckpoint:

# log model only if `val_accuracy` increases
wandb_logger = WandbLogger(log_model="all")
checkpoint_callback = ModelCheckpoint(monitor="val_accuracy", mode="max")
trainer = Trainer(logger=wandb_logger, callbacks=[checkpoint_callback])

latest and best aliases are automatically set to easily retrieve a model checkpoint:

# reference can be retrieved in artifacts panel
# "VERSION" can be a version (ex: "v2") or an alias ("latest or "best")
checkpoint_reference = "USER/PROJECT/MODEL-RUN_ID:VERSION"

# download checkpoint locally (if not already cached)
run = wandb.init(project="MNIST")
artifact = run.use_artifact(checkpoint_reference, type="model")
artifact_dir = artifact.download()

# load checkpoint
model = LitModule.load_from_checkpoint(Path(artifact_dir) / "model.ckpt")

Log media

Log text with:

# using columns and data
columns = ["input", "label", "prediction"]
data = [["cheese", "english", "english"], ["fromage", "french", "spanish"]]
wandb_logger.log_text(key="samples", columns=columns, data=data)

# using a pandas DataFrame
wandb_logger.log_text(key="samples", dataframe=my_dataframe)

Log images with:

# using tensors, numpy arrays or PIL images
wandb_logger.log_image(key="samples", images=[img1, img2])

# adding captions
wandb_logger.log_image(key="samples", images=[img1, img2], caption=["tree", "person"])

# using file path
wandb_logger.log_image(key="samples", images=["img_1.jpg", "img_2.jpg"])

More arguments can be passed for logging segmentation masks and bounding boxes. Refer to Image Overlays documentation.

Log Tables

W&B Tables can be used to log, query and analyze tabular data.

They support any type of media (text, image, video, audio, molecule, html, etc) and are great for storing, understanding and sharing any form of data, from datasets to model predictions.

columns = ["caption", "image", "sound"]
data = [["cheese", wandb.Image(img_1), wandb.Audio(snd_1)], ["wine", wandb.Image(img_2), wandb.Audio(snd_2)]]
wandb_logger.log_table(key="samples", columns=columns, data=data)

Downloading and Using Artifacts

To download an artifact without starting a run, call the download_artifact function on the class:

from lightning.pytorch.loggers import WandbLogger

artifact_dir = WandbLogger.download_artifact(artifact="path/to/artifact")

To download an artifact and link it to an ongoing run call the download_artifact function on the logger instance:

class MyModule(LightningModule):
    def any_lightning_module_function_or_hook(self):
        self.logger.download_artifact(artifact="path/to/artifact")

To link an artifact from a previous run you can use use_artifact function:

from lightning.pytorch.loggers import WandbLogger

wandb_logger = WandbLogger(project="my_project", name="my_run")
wandb_logger.use_artifact(artifact="path/to/artifact")