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, **kwargs)[source]#

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 train, tune
from ray.train import RunConfig
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=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.
**kwargs – The keyword arguments will be pased 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/library/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.