Ray Tune: Hyperparameter tuning#
Ray Tune is a Python library for experiment execution and hyperparameter tuning at any scale. You can tune models built with PyTorch, XGBoost, TensorFlow and Keras, and other machine learning frameworks by running state-of-the-art algorithms such as Population Based Training (PBT) and Hyperband and ASHA. Tune also integrates with many other hyperparameter optimization tools, including Ax, BayesOpt, BOHB, Nevergrad, and Optuna.
Select a tab to see code examples for different machine learning frameworks.
To run this example, install Ray Tune with pip install "ray[tune]".
This quickstart example minimizes the objective function, which computes a**2 + b. The closer a is to zero and the smaller b is, the smaller the result. You define a search space for a and b, and Ray Tune explores the space for good values.
from ray import tune
def objective(config): # ①
score = config["a"] ** 2 + config["b"]
return {"score": score}
search_space = { # ②
"a": tune.grid_search([0.001, 0.01, 0.1, 1.0]),
"b": tune.choice([1, 2, 3]),
}
tuner = tune.Tuner(objective, param_space=search_space) # ③
results = tuner.fit()
print(results.get_best_result(metric="score", mode="min").config)
① Define an objective function.
② Define a search space.
③ Start a Tune run and print the best result.
To tune a Keras model with Hyperopt, wrap the model in an objective function and use the function’s config to select hyperparameters. This example tunes only the activation parameter of the model’s first layer, but you can tune any model parameter. After you define the search space, initialize the HyperOptSearch object and pass it to Tuner as the search_alg argument of tune.TuneConfig. Tell Ray Tune which metric to optimize and whether to maximize or minimize it.
from ray import tune
from ray.tune.search.hyperopt import HyperOptSearch
import keras
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}
search_space = {"activation": tune.choice(["relu", "tanh"])} # ②
algo = HyperOptSearch()
tuner = tune.Tuner( # ③
objective,
tune_config=tune.TuneConfig(
metric="accuracy",
mode="max",
search_alg=algo,
),
param_space=search_space,
)
results = tuner.fit()
① Wrap a Keras model in an objective function.
② Define a search space and initialize the search algorithm.
③ Start a Tune run that maximizes accuracy.
To tune a PyTorch model with Optuna, wrap the model in an objective function and use the function’s config to select hyperparameters. This example tunes only the momentum and lr parameters of the model’s optimizer, where lr is the learning rate, but you can tune any other model parameter. After you define the search space, initialize the OptunaSearch object and pass it to Tuner as the search_alg argument of tune.TuneConfig. Tell Ray Tune which metric to optimize and whether to maximize or minimize it. This example stops each trial after five iterations, but you can define other stopping rules.
import torch
from ray import tune
from ray.tune.search.optuna import OptunaSearch
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_epoch(model, optimizer, train_loader) # Train the model
acc = test(model, test_loader) # Compute test accuracy
tune.report({"mean_accuracy": acc}) # Report to Tune
search_space = {"lr": tune.loguniform(1e-4, 1e-2), "momentum": tune.uniform(0.1, 0.9)}
algo = OptunaSearch() # ②
tuner = tune.Tuner( # ③
objective,
tune_config=tune.TuneConfig(
metric="mean_accuracy",
mode="max",
search_alg=algo,
),
run_config=tune.RunConfig(
stop={"training_iteration": 5},
),
param_space=search_space,
)
results = tuner.fit()
print("Best config is:", results.get_best_result().config)
① Wrap a PyTorch model in an objective function.
② Define a search space and initialize the search algorithm.
③ Start a Tune run that maximizes mean accuracy and stops after five iterations.
With Tune, you can also launch a multi-node distributed hyperparameter sweep in fewer than 10 lines of code. You can then move your models from training to serving on the same infrastructure with Ray Serve.
Getting started
Learn how to tune a PyTorch model with Tune in the getting started tutorial.
Key concepts
Understand the key concepts behind Ray Tune. Learn about trials, search algorithms, schedulers, and other features.
User guides
The user guides cover key features of Tune, such as distributed training and early stopping.
Examples
The examples include practical tutorials for frameworks such as scikit-learn, Keras, TensorFlow, PyTorch, and MLflow, and for state-of-the-art search algorithm integrations.
Ray Tune FAQ
Find answers to common questions in the FAQ.
Ray Tune API
Get detailed information about the Ray Tune API, including search spaces, algorithms, and training configurations.
Why choose Tune?#
Many other hyperparameter optimization libraries exist. If you’re new to Tune, you’re probably wondering what makes Tune different.
Projects using Tune#
To add a project to this list or to request the removal of a listed project, submit a pull request. The following are some of the popular open source repositories and research projects that use Tune:
Softlearning: Softlearning is a reinforcement learning framework for training maximum entropy policies in continuous domains. It includes the official implementation of the Soft Actor-Critic algorithm.
Flambe: Flambe is a machine learning framework that accelerates research and its path to production. See flambe.ai.
Population Based Augmentation: Population Based Augmentation (PBA) is an 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: Allentune is a hyperparameter search tool for AllenNLP from AllenAI.
machinable: The machinable project is 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.
Learn more about Ray Tune#
The following blog posts and talks cover Ray Tune:
Tune: a Python library for fast hyperparameter tuning at any scale, a blog post
Cutting edge hyperparameter tuning with Ray Tune, a blog post
A Guide to Modern Hyperparameter Optimization, a video of a talk at PyData LA 2019, with slides
Cite Tune#
If Tune helps your academic research, consider citing the Tune paper. The following example shows a BibTeX entry:
@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}
}