Using Aim with Tune#

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:

%pip install aim
%pip install ray[tune]

import numpy as np

import ray
from ray import train, 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.

def train_function(config):
    for _ in range(50):
        loss = config["mean"] + config["sd"] * np.random.randn()
        train.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.

tuner = tune.Tuner(
    train_function,
    run_config=train.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()

2023-02-07 00:04:11,228	INFO worker.py:1544 -- Started a local Ray instance. View the dashboard at http://127.0.0.1:8265 

2023-02-07 00:04:19,366	INFO tune.py:798 -- Total run time: 7.38 seconds (6.85 seconds for the tuning loop).

<ray.tune.result_grid.ResultGrid at 0x137de07c0>

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).

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.

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

--------------------------------------------------------------------------
                Aim UI collects anonymous usage analytics.                
                        Read how to opt-out here:                         
    https://aimstack.readthedocs.io/en/latest/community/telemetry.html    
--------------------------------------------------------------------------
Running Aim UI on repo `<Repo#-5734997863388805469 path=/tmp/ray_results/aim_example/.aim read_only=None>`
Open http://127.0.0.1:43800
Press Ctrl+C to exit
^C

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

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

Tune Aim Logger API#

class ray.tune.logger.aim.AimLoggerCallback(repo: str | None = None, experiment_name: str | None = None, metrics: List[str] | None = None, **aim_run_kwargs)[source]

Aim Logger: logs metrics in Aim format.

Aim is an open-source, self-hosted ML experiment tracking tool. It’s good at tracking lots (thousands) of training runs, and it allows you to compare them with a performant and well-designed UI.

Source: aimhubio/aim

Parameters:

• repo – Aim repository directory or a Repo object that the Run object will log results to. If not provided, a default repo will be set up in the experiment directory (one level above trial directories).

• experiment – Sets the experiment property of each Run object, which is the experiment name associated with it. Can be used later to query runs/sequences. If not provided, the default will be the Tune experiment name set by RunConfig(name=...).

• metrics – List of metric names (out of the metrics reported by Tune) to track in Aim. If no metric are specified, log everything that is reported.

• aim_run_kwargs – Additional arguments that will be passed when creating the individual Run objects for each trial. For the full list of arguments, please see the Aim documentation: https://aimstack.readthedocs.io/en/latest/refs/sdk.html