ray.air.integrations.wandb.WandbLoggerCallback

class ray.air.integrations.wandb.WandbLoggerCallback(project: str | None = None, group: str | None = None, api_key_file: str | None = None, api_key: str | None = None, excludes: List[str] | None = None, log_config: bool = False, upload_checkpoints: bool = False, save_checkpoints: bool = False, upload_timeout: int = 1800, video_kwargs: dict | None = None, image_kwargs: dict | None = None, **kwargs)

Bases: LoggerCallback

Weights and biases (https://www.wandb.ai/) is a tool for experiment tracking, model optimization, and dataset versioning. This Ray Tune LoggerCallback sends metrics to Wandb for automatic tracking and visualization.

Example

import random

from ray import tune
from ray.air.integrations.wandb import WandbLoggerCallback


def train_func(config):
    offset = random.random() / 5
    for epoch in range(2, config["epochs"]):
        acc = 1 - (2 + config["lr"]) ** -epoch - random.random() / epoch - offset
        loss = (2 + config["lr"]) ** -epoch + random.random() / epoch + offset
        train.report({"acc": acc, "loss": loss})


tuner = tune.Tuner(
    train_func,
    param_space={
        "lr": tune.grid_search([0.001, 0.01, 0.1, 1.0]),
        "epochs": 10,
    },
    run_config=tune.RunConfig(
        callbacks=[WandbLoggerCallback(project="Optimization_Project")]
    ),
)
results = tuner.fit()

Parameters:

• project – Name of the Wandb project. Mandatory.

• group – Name of the Wandb group. Defaults to the trainable name.

• api_key_file – Path to file containing the Wandb API KEY. This file only needs to be present on the node running the Tune script if using the WandbLogger.

• api_key – Wandb API Key. Alternative to setting api_key_file.

• excludes – List of metrics and config that should be excluded from the log.

• log_config – Boolean indicating if the config parameter of the results dict should be logged. This makes sense if parameters will change during training, e.g. with PopulationBasedTraining. Defaults to False.

• upload_checkpoints – If True, model checkpoints will be uploaded to Wandb as artifacts. Defaults to False.

• video_kwargs – Dictionary of keyword arguments passed to wandb.Video() when logging videos. Videos have to be logged as 5D numpy arrays to be affected by this parameter. For valid keyword arguments, see https://docs.wandb.ai/ref/python/data-types/video/. Defaults to None.

• image_kwargs – Dictionary of keyword arguments passed to wandb.Image() when logging images. Images have to be logged as 3D or 4D numpy arrays to be affected by this parameter. For valid keyword arguments, see https://docs.wandb.ai/ref/python/data-types/image/. Defaults to None.

• **kwargs – The keyword arguments will be passed to wandb.init().

Wandb’s group, run_id and run_name are automatically selected by Tune, but can be overwritten by filling out the respective configuration values.

Please see here for all other valid configuration settings: https://docs.wandb.ai/ref/python/init/

PublicAPI (alpha): This API is in alpha and may change before becoming stable.

Methods

get_state	Get the state of the callback.
log_trial_restore	Handle logging when a trial restores.
on_checkpoint	Called after a trial saved a checkpoint with Tune.
on_experiment_end	Wait for the actors to finish their call to wandb.finish.
on_step_begin	Called at the start of each tuning loop step.
on_step_end	Called at the end of each tuning loop step.
on_trial_recover	Called after a trial instance failed (errored) but the trial is scheduled for retry.
set_state	Set the state of the callback.

Attributes

AUTO_CONFIG_KEYS

Results that are saved with wandb.config instead of wandb.log.