Ray Tune: a Python library for fast hyperparameter tuning at any scale | by Richard Liaw | Towards Data Science

If you’ve ever tried to tune hyperparameters for a machine learning model, you know that it can be a very painful process. Simple approaches quickly become time-consuming.
And now more than ever, you absolutely need cutting-edge hyperparameter tuning tools to keep up with the state-of-the-art.
Model advancements are becoming more and more dependent on newer and better hyperparameter tuning algorithms such as Population Based Training (PBT), HyperBand, and ASHA.
These algorithms provide two critical benefits:
They maximize model performance: e.g., DeepMind uses PBT to achieve superhuman performance on StarCraft; Waymo uses PBT to enable self-driving cars.
They minimize training costs: HyperBand and ASHA converge to high-quality configurations in half the time taken by previous approaches; population-based data augmentation algorithms cuts costs by orders of magnitude.
The fact of the matter is that the vast majority of researchers and teams do not leverage such algorithms. Most existing hyperparameter search frameworks do not have these newer optimization algorithms. And once you reach a certain scale, most existing solutions for parallel hyperparameter search can be a hassle to use — you’ll need to configure each machine for each run and often manage a separate database.
Practically speaking, implementing and maintaining these algorithms requires a significant amount of time and engineering.
But it doesn’t need to be this way. There’s no reason why you can’t easily incorporate hyperparameter tuning into your machine learning project, seamlessly run a parallel asynchronous grid search on 8 GPUs in your cluster, and leverage Population Based Training or any Bayesian optimization algorithm at scale on the cloud.
In this blog post, we’ll introduce RayTune, a powerful hyperparameter optimization library designed to remove the friction from scaling experiment execution and hyperparameter search.
RayTune provides distributed asynchronous optimization out of the box.
RayTune offers state of the art algorithms including (but not limited to) ASHA, BOHB, and Population-Based Training.
You can scale a RayTune hyperparameter search from a single machine to a large distributed cluster without changing your code.
RayTune integrates with many optimization libraries such as Ax/Botorch, HyperOpt, and Bayesian Optimization and enables you to scale them transparently.
RayTune supports any machine learning framework, including PyTorch, TensorFlow, XGBoost, LightGBM, scikit-learn, and Keras.
Beyond RayTune’s core features, there are two primary reasons why researchers and developers prefer RayTune over other existing hyperparameter tuning frameworks: scale and flexibility.
Leverage all of the cores and GPUs on your machine to perform parallel asynchronous hyperparameter tuning by adding fewer than 10 lines of Python.
With another configuration file and 4 lines of code, launch a massive distributed hyperparameter search on the cloud and automatically shut down the machines (we’ll show you how to do this below).
With Tune’s built-in fault tolerance, trial migration, and cluster autoscaling, you can safely leverage spot (preemptible) instances and reduce cloud costs by up to 90%.
Tune integrates seamlessly with experiment management tools such as MLFlow and TensorBoard.
Tune provides a flexible interface for optimization algorithms, allowing you to easily implement and scale new optimization algorithms.
You can use Tune to leverage and scale many state-of-the-art search algorithms and libraries such as HyperOpt (below) and Ax without modifying any model training code.
Let’s now dive into a concrete example that shows how you to leverage a state-of-the-art early stopping algorithm (ASHA). We will start by running Tune across all of the cores on your workstation. We’ll then scale out the same experiment on the cloud with about 10 lines of code.
We’ll be using PyTorch in this example, but we also have examples for Tensorflow and Keras available.
Tune is installed as part of Ray. To run this example, you will need to install the following: pip install ray torch torchvision.
You can download a full version of the blog in this blog here.
We first run some imports :
Let’s write a neural network with PyTorch:
To start using Tune, add a simple logging statement to the PyTorch training below function.
def train_mnist(config):    model = ConvNet()    train_loader, test_loader = get_data_loaders()    optimizer = optim.SGD(        model.parameters(),         lr=config["lr"],         momentum=config["momentum"])    for i in range(10):        train(model, optimizer, train_loader, torch.device("cpu"))        acc = test(model, test_loader, torch.device("cpu"))        tune.track.log(mean_accuracy=acc)        if i % 5 == 0:            # This saves the model to the trial directory            torch.save(model, "./model.pth")
Notice that there’s a couple helper functions in the above training script; you can see their definitions here.
Let’s run 1 trial, randomly sampling from a uniform distribution for learning rate and momentum.
Now, you’ve run your first Tune run! You can easily enable GPU usage by specifying GPU resources — see the documentation for more details. We can then plot the performance of this trial.
Let’s integrate ASHA, a scalable algorithm for early stopping (blog post and paper). ASHA terminates trials that are less promising and allocates more time and resources to more promising trials.
Parallelize your search across all available cores on your machine with num_samples (extra trials will be queued).
You can use the same DataFrame plotting as the previous example. After running, if Tensorboard is installed, you can also use Tensorboard for visualizing results: tensorboard --logdir ~/ray_results
Setting up a distributed hyperparameter search is often too much work. Tune and Ray make this seamless.
First, we’ll create a YAML file which configures a Ray cluster. As part of Ray, Tune interoperates very cleanly with the Ray cluster launcher. The same commands shown below will work on GCP, AWS, and local private clusters. We’ll use 3 worker nodes in addition to a head node, so we should have a total of 32 vCPUs on the cluster — allowing us to evaluate 32 hyperparameter configurations in parallel.
To distribute your hyperparameter search across the Ray cluster, you’ll need to append this to the top of your script:
Given the large increase in compute, we should be able to increase our search space and number of samples in our search space:
You can download a full version of the script in this blog here (as tune_script.py).
To launch your experiment, you can run (assuming your code so far is in a file tune_script.py):
$ ray submit tune-default.yaml tune_script.py --start \     --args=”localhost:6379”
This will launch your cluster on AWS, upload tune_script.py onto the head node, and run python tune_script localhost:6379, which is a port opened by Ray to enable distributed execution.
All of the output of your script will show up on your console. Note that the cluster will setup the head node first before any of the worker nodes, so at first you may see only 4 CPUs available. After some time, you can see 24 trials being executed in parallel, and the other trials will be queued up to be executed as soon as a trial is free.
To shut down your cluster, you can run:
$ ray down tune-default.yaml
And you’re done 🎉!
Tune has numerous other features that enable researchers and practitioners to accelerate their development. Other Tune features not covered in this blogpost include:
A simple API for running distributed fault-tolerant experiments
Distributed Hyperparameter search over Distributed Data Parallel training for PyTorch
Population-based Training
For users that have access to the cloud, Tune and Ray provide a number of utilities that enable a seamless transition between development on your laptop and execution on the cloud. The documentation includes:
running the experiment in a background session
submitting trials to an existing experiment
visualizing all results of a distributed experiment in TensorBoard.
Tune is designed to scale experiment execution and hyperparameter search with ease. If you have any comments or suggestions or are interested in contributing to Tune, you can reach out to me or the ray-dev mailing list.
Code: https://github.com/ray-project/ray/tree/master/python/ray/tuneDocs: http://ray.readthedocs.io/en/latest/tune.html
For other readings on hyperparameter tuning, check out Neptune.ai’s blog post on Optuna vs HyperOpt!
Thanks to Allan Peng, Eric Liang, Joey Gonzalez, Ion Stoica, Eugene Vinitsky, Lisa Dunlap, Philipp Moritz, Andrew Tan, Alvin Wan, Daniel Rothchild, Brijen Thananjeyan, Alok Singh (and maybe others?) for reading through various versions of this blog post!