Machine Learning Toolbox Review: Ax

The goal of the Machine Learning Toolbox Review is to identify tools that accelerate the Machine Learning Process based on the 3 key factors.

1. Onboarding Experience
2. The Value Add
3. Community Support

Each episode will rate how well the tool does and provides an overall grade for the tool.

What is Ax?

Ax is an accessible, general-purpose platform for understanding, managing, deploying, and automating adaptive experiments.

Basically given a set of all possible solutions, find the optimal solution for what you are trying to solve.

Ax was open sourced by Facebook and was released May 1 at their F8 Conference.

Onboarding Experience

The onboarding experience evaluates how seamless it is for a developer to begin using the product, starting from the product page.

Downloading

Downloading is straight forward, using pip3 (sorry python 2.7 folks, but it's time to move on)

pip3 install ax-platform

After the lengthy install, there is sample code given to try to make sure it works

from ax import optimize
best_parameters, best_values, experiment, model = optimize(
        parameters=[
          {
            "name": "x1",
            "type": "range",
            "bounds": [-10.0, 10.0],
          },
          {
            "name": "x2",
            "type": "range",
            "bounds": [-10.0, 10.0],
          },
        ],
        # Booth function
        evaluation_function=lambda p: (p["x1"] + 2*p["x2"] - 7)**2 + (2*p["x1"] + p["x2"] - 5)**2,
        minimize=True,
    )

This sample code wants to find the smallest numbers that can satisfy the equation. The experiment is immediately executed, which takes a couple of seconds. The results are then stored in the best_paramters variable.

{'x1': 0.7810791992092341, 'x2': 3.0911807098099935}

Pretty simple to use, and there are tutorials provided on the site to explore more use cases. https://ax.dev/tutorials/

With no hiccups in the install process, tutorials to get started, and plenty of documentation, this was a great onboarding experience, I give it 5/5.

The Value Add

We can define value add, as how much time will using this tool save me. It could be doing more in fewer lines of code or getting rid of manual processes.

Ax comes with lots of pre-built high level APIs that encapsulate lots of dense machine learning code, to a simple call. For example to do hyperparameter training we would use

def train_evaluate(parameterization):
    net = train(train_loader=train_loader, parameters=parameterization, dtype=dtype, device=device)
    return evaluate(
        net=net,
        data_loader=valid_loader,
        dtype=dtype,
        device=device,
    )

best_parameters, values, experiment, model = optimize(
    parameters=[
        {"name": "lr", "type": "range", "bounds": [1e-6, 0.4], "log_scale": True},
        {"name": "momentum", "type": "range", "bounds": [0.0, 1.0]},
    ],
    evaluation_function=train_evaluate,
    objective_name='accuracy',
)

Classic methods make you have to define each step in a very verbose manner.

Ax also comes with notebook plotting tools make it work seamless with Jupyter notebooks.

The extensive amounts of options, and the encapsulation of doing adaptive experiments gives it a great value add for anyone in this domain, making it a clear 5/5.

Community Support

Community Support can be defined as the community around helping each other solve problems with the tool. This can be through Stackoverflow, GitHub, or any other type of social channel.

Well since Ax is new, there is not no chatter behind it. Nothing on Stackoverflow, and only a handful of issues on the Github page. The contributors are being responsive as the tool was recently released, but with not much of a community built around the tool, I can only recommend using if your an experienced developer in this realm, that is able to accomplish much without support. My grade 3/5

Overall

With a great onboarding experience, and great value add demonstrated by its robust APIs, this tool will be a major player in adaptive experiments. The lack of support for Ax makes it hard to break in for beginners. Hopefully, a community develops around Ax paving a pathway to make it more approachable. Final Rating 4/5