Tune: Scalable Hyperparameter Tuning

../_images/tune_overview.png

Tune is a Python library for experiment execution and hyperparameter tuning at any scale. You can tune your favorite machine learning framework, including PyTorch, XGBoost, TensorFlow and Keras, and choose among state of the art algorithms such as Population Based Training (PBT), BayesOptSearch, or HyperBand/ASHA. Tune integrates with a wide range of hyperparameter optimization tools, like Optuna, Hyperopt, Ax, and Nevergrad, to name a few.

Learn how to use Ray Tune for various machine learning frameworks in just a few steps. Click on the tabs to see code examples.

Tip

We’d love to hear your feedback on using Tune - get in touch!

To run this example, install the following: pip install "ray[tune]".

In this quick-start example you minimize a simple function of the form f(x) = a**2 + b, our objective function. The closer a is to zero and the smaller b is, the smaller the total value of f(x). We will define a so-called search space for a and b and let Ray Tune explore the space for good values.

from ray import tune

# 1. Define an objective function.
def objective(config):
    score = config["a"] ** 2 + config["b"]
    return {"score": score}


# 2. Define a search space.
search_space = {
    "a": tune.grid_search([0.001, 0.01, 0.1, 1.0]),
    "b": tune.choice([1, 2, 3]),
}

# 3. Start a Tune run and print the best result.
analysis = tune.run(objective, config=search_space)
print(analysis.get_best_config(metric="score", mode="min"))

To tune your Keras models with Hyperopt, you wrap your model in an objective function whose config you can access for selecting hyperparameters. In the example below we only tune the activation parameter of the first layer of the model, but you can tune any parameter of the model you want. After defining the search space, you can simply initialize the HyperOptSearch object and pass it to run. It’s important to tell Ray Tune which metric you want to optimize and whether you want to maximize or minimize it.

from ray import tune
from ray.tune.suggest.hyperopt import HyperOptSearch
import keras


# 1. Wrap a Keras model in an objective function.
def objective(config):
    model = keras.models.Sequential()
    model.add(keras.layers.Dense(784, activation=config["activation"]))
    model.add(keras.layers.Dense(10, activation="softmax"))

    model.compile(loss="binary_crossentropy", optimizer="adam", metrics=["accuracy"])
    # model.fit(...)
    # loss, accuracy = model.evaluate(...)
    return {"accuracy": accuracy}


# 2. Define a search space and initialize the search algorithm.
search_space = {"activation": tune.choice(["relu", "tanh"])}
algo = HyperOptSearch()

# 3. Start a Tune run that maximizes accuracy.
analysis = tune.run(
    objective, search_alg=algo, config=search_space, metric="accuracy", mode="max"
)

To tune your PyTorch models with Optuna, you wrap your model in an objective function whose config you can access for selecting hyperparameters. In the example below we only tune the momentum and learning rate (lr) parameters of the model’s optimizer, but you can tune any other model parameter you want. After defining the search space, you can simply initialize the OptunaSearch object and pass it to run. It’s important to tell Ray Tune which metric you want to optimize and whether you want to maximize or minimize it. We stop tuning this training run after 5 iterations, but you can easily define other stopping rules as well.

# 1. Wrap your PyTorch model in an objective function.
import torch
from ray import tune
from ray.tune.suggest.optuna import OptunaSearch


# 1. Wrap a PyTorch model in an objective function.
def objective(config):
    train_loader, test_loader = load_data()  # Load some data
    model = ConvNet().to("cpu")  # Create a PyTorch conv net
    optimizer = torch.optim.SGD(  # Tune the optimizer
        model.parameters(), lr=config["lr"], momentum=config["momentum"]
    )

    while True:
        train(model, optimizer, train_loader)  # Train the model
        acc = test(model, test_loader)  # Compute test accuracy
        tune.report(mean_accuracy=acc)  # Report to Tune


# 2. Define a search space and initialize the search algorithm.
search_space = {"lr": tune.loguniform(1e-4, 1e-2), "momentum": tune.uniform(0.1, 0.9)}
algo = OptunaSearch()

# 3. Start a Tune run that maximizes mean accuracy and stops after 5 iterations.
analysis = tune.run(
    objective,
    metric="mean_accuracy",
    mode="max",
    search_alg=algo,
    stop={"training_iteration": 5},
    config=search_space,
)
print("Best config is:", analysis.best_config)

With Tune you can also launch a multi-node distributed hyperparameter sweep in less than 10 lines of code. It automatically manages checkpoints and logging to TensorBoard. And you can move your models from training to serving on the same infrastructure with Ray Serve.

Getting Started

In our getting started tutorial you will learn how to tune a PyTorch model effectively with Tune.

Key Concepts

Understand the key concepts behind Ray Tune. Learn about tune runs, search algorithms, schedulers and other features.

User Guides

Our guides teach you about key features of Tune, such as distributed training or early stopping.

Examples

In our examples you can find practical tutorials for scikit-learn, Keras, TensorFlow, PyTorch, mlflow, and many more.

Ray Tune FAQ

Find answers to commonly asked questions in our detailed FAQ.

Ray Tune API

Get more in-depth information about the Ray Tune API, including all about search spaces, algorithms and training configurations.

Why choose Tune?

There are many other hyperparameter optimization libraries out there. If you’re new to Tune, you’re probably wondering, “what makes Tune different?”

Projects using Tune

Here are some of the popular open source repositories and research projects that leverage Tune. Feel free to submit a pull-request adding (or requesting a removal!) of a listed project.

  • Softlearning: Softlearning is a reinforcement learning framework for training maximum entropy policies in continuous domains. Includes the official implementation of the Soft Actor-Critic algorithm.

  • Flambe: An ML framework to accelerate research and its path to production. See flambe.ai.

  • Population Based Augmentation: Population Based Augmentation (PBA) is a algorithm that quickly and efficiently learns data augmentation functions for neural network training. PBA matches state-of-the-art results on CIFAR with one thousand times less compute.

  • Fast AutoAugment by Kakao: Fast AutoAugment (Accepted at NeurIPS 2019) learns augmentation policies using a more efficient search strategy based on density matching.

  • Allentune: Hyperparameter Search for AllenNLP from AllenAI.

  • machinable: A modular configuration system for machine learning research. See machinable.org.

  • NeuroCard: NeuroCard (Accepted at VLDB 2021) is a neural cardinality estimator for multi-table join queries. It uses state of the art deep density models to learn correlations across relational database tables.

Citing Tune

If Tune helps you in your academic research, you are encouraged to cite our paper. Here is an example bibtex:

@article{liaw2018tune,
    title={Tune: A Research Platform for Distributed Model Selection and Training},
    author={Liaw, Richard and Liang, Eric and Nishihara, Robert
            and Moritz, Philipp and Gonzalez, Joseph E and Stoica, Ion},
    journal={arXiv preprint arXiv:1807.05118},
    year={2018}
}