Get started with Ray Tune#

This tutorial walks you through setting up a Ray Tune experiment. Starting from a PyTorch model, you use Tune to optimize the model’s hyperparameters with early stopping and Bayesian optimization through Hyperopt.

Tip

If you have suggestions on how to improve this tutorial, open an issue on GitHub.

To run this example, install the following packages:

pip install "ray[tune]" torch torchvision

Set up a PyTorch model to tune#

First, import the dependencies. The PyTorch and torchvision modules help you create and train a model, and Ray Tune helps you optimize it. The imports also include a scheduler, in this case the ASHAScheduler, which you use to tune the model later in this tutorial.

import numpy as np
import torch
import torch.optim as optim
import torch.nn as nn
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torch.nn.functional as F

from ray import tune
from ray.tune.schedulers import ASHAScheduler

Next, define a simple PyTorch model to train. If you’re not familiar with PyTorch, the simplest way to define a model is to implement an nn.Module. Set up the model in __init__, and then implement the forward pass. This example uses a small convolutional neural network with one 2D convolutional layer, a fully connected layer, and a softmax function.

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        # In this example, we don't change the model architecture
        # due to simplicity.
        self.conv1 = nn.Conv2d(1, 3, kernel_size=3)
        self.fc = nn.Linear(192, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 3))
        x = x.view(-1, 192)
        x = self.fc(x)
        return F.log_softmax(x, dim=1)

The following train and test functions train and evaluate your PyTorch model. If you already know how to write them, skip ahead to the next section.

Set up a `Tuner` for a training run with Tune#

The following function trains the PyTorch model for multiple epochs. This function runs in a separate process as a Ray actor, so you need to report the model’s performance back to Tune, which runs in the main Python process.

To do this, call tune.report() in the training function, which sends the performance value back to Tune. Because the function runs in a separate process, make sure it’s serializable by Ray.

import os
import tempfile

from ray.tune import Checkpoint

def train_mnist(config):
    # Data Setup
    mnist_transforms = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.1307, ), (0.3081, ))])

    train_loader = DataLoader(
        datasets.MNIST("~/data", train=True, download=True, transform=mnist_transforms),
        batch_size=64,
        shuffle=True)
    test_loader = DataLoader(
        datasets.MNIST("~/data", train=False, transform=mnist_transforms),
        batch_size=64,
        shuffle=True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = ConvNet()
    model.to(device)

    optimizer = optim.SGD(
        model.parameters(), lr=config["lr"], momentum=config["momentum"])
    for i in range(10):
        train_func(model, optimizer, train_loader)
        acc = test_func(model, test_loader)


        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
            checkpoint = None
            if (i + 1) % 5 == 0:
                # This saves the model to the trial directory
                torch.save(
                    model.state_dict(),
                    os.path.join(temp_checkpoint_dir, "model.pth")
                )
                checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)

            # Send the current training result back to Tune
            tune.report({"mean_accuracy": acc}, checkpoint=checkpoint)

Run one trial by calling Tuner.fit, and randomly sample the learning rate from a log-uniform distribution and the momentum from a uniform distribution.

search_space = {
    "lr": tune.sample_from(lambda spec: 10 ** (-10 * np.random.rand())),
    "momentum": tune.uniform(0.1, 0.9),
}

# Uncomment this to enable distributed execution
# `ray.init(address="auto")`

# Download the dataset first
datasets.MNIST("~/data", train=True, download=True)

tuner = tune.Tuner(
    train_mnist,
    param_space=search_space,
)
results = tuner.fit()

Tuner.fit returns a ResultGrid object. You can use it to plot the trial’s performance.

dfs = {result.path: result.metrics_dataframe for result in results}
[d.mean_accuracy.plot() for d in dfs.values()]

Note

Tune automatically runs trials in parallel across all available cores and GPUs on your machine or cluster. To limit the number of concurrent trials, use the ConcurrencyLimiter.

Stop trials early with Asynchronous Successive Halving#

Next, add early stopping to the optimization process with Asynchronous Successive Halving (ASHA), a scalable algorithm for principled early stopping.

At a high level, ASHA stops less promising trials and allocates more time and resources to more promising ones. Because the optimization process becomes more efficient, you can afford to run 5x more trials by adjusting the num_samples parameter.

Tune implements ASHA as a trial scheduler, ASHAScheduler. Trial schedulers can stop bad trials early, pause trials, clone trials, and alter the hyperparameters of a running trial. For details on the available schedulers and library integrations, see Tune trial schedulers.

tuner = tune.Tuner(
    train_mnist,
    tune_config=tune.TuneConfig(
        num_samples=20,
        scheduler=ASHAScheduler(metric="mean_accuracy", mode="max"),
    ),
    param_space=search_space,
)
results = tuner.fit()

# Obtain a trial dataframe from all run trials of this `tune.run` call.
dfs = {result.path: result.metrics_dataframe for result in results}

To visualize trial progress, you can run the following code in a Jupyter notebook.

# Plot by epoch
ax = None  # This plots everything on the same plot
for d in dfs.values():
    ax = d.mean_accuracy.plot(ax=ax, legend=False)

You can also visualize results with TensorBoard.

tensorboard --logdir {logdir}

Use search algorithms in Tune#

In addition to trial schedulers, you can further optimize your hyperparameters with an intelligent search technique such as Bayesian optimization. To do this, you can use a Tune search algorithm. Search algorithms use optimization algorithms to navigate the given hyperparameter space intelligently.

Each search library has its own way of defining the search space. This example uses Hyperopt’s hp functions with HyperOptSearch.

from hyperopt import hp
from ray.tune.search.hyperopt import HyperOptSearch

space = {
    "lr": hp.loguniform("lr", -10, -1),
    "momentum": hp.uniform("momentum", 0.1, 0.9),
}

hyperopt_search = HyperOptSearch(space, metric="mean_accuracy", mode="max")

tuner = tune.Tuner(
    train_mnist,
    tune_config=tune.TuneConfig(
        num_samples=10,
        search_alg=hyperopt_search,
    ),
)
results = tuner.fit()

# To enable GPUs, use this instead:
# analysis = tune.run(
#     train_mnist, config=search_space, resources_per_trial={'gpu': 1})

Note

You can combine some search algorithms with different trial schedulers. For details, see Tune trial schedulers.

Evaluate your model after tuning#

To evaluate the best trained model, retrieve it with the ResultGrid object:

best_result = results.get_best_result("mean_accuracy", mode="max")
with best_result.checkpoint.as_directory() as checkpoint_dir:
    state_dict = torch.load(os.path.join(checkpoint_dir, "model.pth"))

model = ConvNet()
model.load_state_dict(state_dict)

Next steps#

See the Tune tutorials for guides on using Tune with your preferred machine learning library.
Browse the gallery of examples to see how to use Tune with PyTorch, XGBoost, TensorFlow, and other libraries.
If you run into problems or have questions, open an issue on GitHub.
To check on your application, you can use the Ray dashboard.