Back to Ray

Using Aim with Tune

doc/source/tune/examples/tune-aim.ipynb

1.13.13.9 KB
Original Source

(tune-aim-ref)=

Using Aim with Tune

<a id="try-anyscale-quickstart-tune-aim" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=tune-aim"> </a> </br>

Aim is an easy-to-use and supercharged open-source experiment tracker. Aim logs your training runs, enables a well-designed UI to compare them, and provides an API to query them programmatically.

{image}
:align: center
:alt: Aim
:width: 100%
:target: https://aimstack.io

Ray Tune currently offers built-in integration with Aim. The {ref}AimLoggerCallback <tune-aim-logger> automatically logs metrics that are reported to Tune by using the Aim API.

{contents}
:backlinks: none
:local: true

Logging Tune Hyperparameter Configurations and Results to Aim

The following example demonstrates how the AimLoggerCallback can be used in a Tune experiment. Begin by installing and importing the necessary modules:

python
%pip install aim
%pip install ray[tune]
python
import numpy as np

import ray
from ray import tune
from ray.tune.logger.aim import AimLoggerCallback

Next, define a simple train_function, which is a Trainable that reports a loss to Tune. The objective function itself is not important for this example, as our main focus is on the integration with Aim.

python
def train_function(config):
    for _ in range(50):
        loss = config["mean"] + config["sd"] * np.random.randn()
        tune.report({"loss": loss})

Here is an example of how you can use the AimLoggerCallback with simple grid-search Tune experiment. The logger will log each of the 9 grid-search trials as separate Aim runs.

python
tuner = tune.Tuner(
    train_function,
    run_config=tune.RunConfig(
        callbacks=[AimLoggerCallback()],
        storage_path="/tmp/ray_results",
        name="aim_example",
    ),
    param_space={
        "mean": tune.grid_search([1, 2, 3, 4, 5, 6, 7, 8, 9]),
        "sd": tune.uniform(0.1, 0.9),
    },
    tune_config=tune.TuneConfig(
        metric="loss",
        mode="min",
    ),
)
tuner.fit()

When the script executes, a grid-search is carried out and the results are saved to the Aim repo, stored at the default location -- the experiment log directory (in this case, it's at /tmp/ray_results/aim_example).

More Configuration Options for Aim

In the example above, we used the default configuration for the AimLoggerCallback. There are a few options that can be configured as arguments to the callback. For example, setting AimLoggerCallback(repo="/path/to/repo") will log results to the Aim repo at that filepath, which could be useful if you have a central location where the results of multiple Tune experiments are stored. Relative paths to the working directory where Tune script is launched can be used as well. By default, the repo will be set to the experiment log directory. See the API reference for more configurations.

Launching the Aim UI

Now that we have logged our results to the Aim repository, we can view it in Aim's web UI. To do this, we first find the directory where the Aim repository lives, then we use the Aim CLI to launch the web interface.

python
# Uncomment the following line to launch the Aim UI!
#!aim up --repo=/tmp/ray_results/aim_example

After launching the Aim UI, we can open the web interface at localhost:43800.

{image}
:align: center
:alt: Aim Metrics Explorer
:target: https://aimstack.readthedocs.io/en/latest/ui/pages/explorers.html#metrics-explorer

The next sections contain more in-depth information on the API of the Tune-Aim integration.

Tune Aim Logger API

(tune-aim-logger)=

{eval-rst}
.. autoclass:: ray.tune.logger.aim.AimLoggerCallback
   :noindex: