Defining metrics for Weights & Biases (wandb)

Ray Libraries (Data, Train, Tune, Serve)
1

TL;DR: How can we call wandb.define_metrics from within ray tune?

Context: The ray.air.integrations.wandb.WandbLoggerCallback allows to loc metrics to wandb after every epoch. However, it does not deal with summary statistics which can be defined as follows:

wandb.init()
wandb.define_metric("accuracy", summary="max")

This tells wandb that the summary metric for accuracy should be the maximal value from every epoch.

What is the best idea to replicate this with ray tune? I am only aware of hacky solutions (sneaking it in the trainable setup – but then wandb.init() will be called twice) etc.

Let me know if I should rather open a feature request for this.

2

Hi @kemok,

This is currently not possible with the WandbLoggerCallback, definitely create a feature request on Github if you have the chance! We can add a custom init function that allows you to define these summary metrics.

For now, if you use the setup_wandb method of integrating Tune with W&B instead, then you have full customizability of what you do with wandb. The only downside here is that you have to log metrics manually, rather than all the Tune metrics getting logged automatically.

For example:

def train_function_wandb(config):
    wandb = setup_wandb(config, api_key="<your-api-key>")
    wandb.define_metric(...)

    for i in range(30):
        loss = config["mean"] + config["sd"] * np.random.randn()
        session.report({"loss": loss})
        wandb.log(dict(loss=loss))

See here for a full example: Using Weights & Biases with Tune — Ray 2.4.0