Ray Tune
One is the
WandbLogger
, which automatically logs metrics reported to Tune to the Wandb API. The other one is the @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.from ray.tune.integration.wandb import WandbLogger
Wandb configuration is done by passing a wandb key to the config parameter of
tune.run()
(see example below).The content of the wandb config entry is passed to
wandb.init()
as keyword arguments. The exception are the following settings, which are used to configure the WandbLogger
itself:api_key_file (str)
– Path to file containing the Wandb API KEY
.api_key (str)
– Wandb API Key. Alternative to setting api_key_file
.excludes (list)
– List of metrics that should be excluded from the log
.log_config (bool)
– 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.from ray.tune.logger import DEFAULT_LOGGERS
from ray.tune.integration.wandb import WandbLogger
tune.run(
train_fn,
config={
# define search space here
"parameter_1": tune.choice([1, 2, 3]),
"parameter_2": tune.choice([4, 5, 6]),
# wandb configuration
"wandb": {
"project": "Optimization_Project",
"api_key_file": "/path/to/file",
"log_config": True
}
},
loggers=DEFAULT_LOGGERS + (WandbLogger, ))
ray.tune.integration.wandb.wandb_mixin(func)
This Ray Tune Trainable
mixin
helps initializing the Wandb API for use with the Trainable
class or with @wandb_mixin
for the function API.For basic usage, just prepend your training function with the
@wandb_mixin
decorator:from ray.tune.integration.wandb import wandb_mixin
@wandb_mixin
def train_fn(config):
wandb.log()
Wandb configuration is done by passing a
wandb key
to the config
parameter of tune.run()
(see example below).The content of the wandb config entry is passed to
wandb.init()
as keyword arguments. The exception are the following settings, which are used to configure the WandbTrainableMixin
itself:api_key_file (str)
– Path to file containing the Wandb API KEY
.api_key (str)
– Wandb API Key. Alternative to setting api_key_file
.Wandb’s
group
, run_id
and run_name
are automatically selected by Tune, but can be overwritten by filling out the respective configuration values.from ray import tune
from ray.tune.integration.wandb import wandb_mixin
@wandb_mixin
def train_fn(config):
for i in range(10):
loss = self.config["a"] + self.config["b"]
wandb.log({"loss": loss})
tune.report(loss=loss)
tune.run(
train_fn,
config={
# define search space here
"a": tune.choice([1, 2, 3]),
"b": tune.choice([4, 5, 6]),
# wandb configuration
"wandb": {
"project": "Optimization_Project",
"api_key_file": "/path/to/file"
}
})
We've created a few examples for you to see how the integration works:
Last modified 1yr ago