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)[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 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 theresults
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 toFalse
.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
andrun_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 the state of the callback.
Handle logging when a trial restores.
Called after a trial saved a checkpoint with Tune.
Wait for the actors to finish their call to
wandb.finish
.Called at the start of each tuning loop step.
Called at the end of each tuning loop step.
Called after a trial instance failed (errored) but the trial is scheduled for retry.
Set the state of the callback.
Attributes
Results that are saved with
wandb.config
instead ofwandb.log
.