Using Weights & Biases with Tune

Weights & Biases (Wandb) is a tool for experiment tracking, model optimizaton, and dataset versioning. It is very popular in the machine learning and data science community for its superb visualization tools.

Ray Tune currently offers two lightweight integrations for Weights & Biases. One is the {ref}WandbLoggerCallback <tune-wandb-logger>, which automatically logs metrics reported to Tune to the Wandb API.

The other one is the {ref}@wandb_mixin <tune-wandb-mixin> decorator, which can be used with the function API. It automatically initializes the Wandb API with Tune's training information. You can just use the Wandb API like you would normally do, e.g. using wandb.log() to log your training process.

Running A Weights & Biases Example

In the following example we're going to use both of the above methods, namely the WandbLoggerCallback and the wandb_mixin decorator to log metrics. Let's start with a few crucial imports:

Next, let's define an easy objective function (a Tune Trainable) that reports a random loss to Tune. The objective function itself is not important for this example, since we want to focus on the Weights & Biases integration primarily.

Given that you provide an api_key_file pointing to your Weights & Biases API key, you cna define a simple grid-search Tune run using the WandbLoggerCallback as follows:

To use the wandb_mixin decorator, you can simply decorate the objective function from earlier. Note that we also use wandb.log(...) to log the loss to Weights & Biases as a dictionary. Otherwise, the decorated version of our objective is identical to its original.

With the decorated_objective defined, running a Tune experiment is as simple as providing this objective and passing the api_key_file to the wandb key of your Tune config:

Finally, you can also define a class-based Tune Trainable by using the WandbTrainableMixin to define your objective:

Running Tune with this WandbTrainable works exactly the same as with the function API. The below tune_trainable function differs from tune_decorated above only in the first argument we pass to Tuner():

Since you may not have an API key for Wandb, we can mock the Wandb logger and test all three of our training functions as follows. If you do have an API key file, make sure to set mock_api to False and pass in the right api_key_file below.

This completes our Tune and Wandb walk-through. In the following sections you can find more details on the API of the Tune-Wandb integration.

Tune Wandb API Reference

WandbLoggerCallback

(tune-wandb-logger)=

.. autoclass:: ray.air.callbacks.wandb.WandbLoggerCallback
   :noindex:

Wandb-Mixin

(tune-wandb-mixin)=

.. autofunction:: ray.tune.integration.wandb.wandb_mixin
   :noindex: