Ray Tune 入门#

本教程将引导您完成设置 Tune 实验的过程。为了入门，我们以一个 PyTorch 模型为例，向您展示如何利用 Ray Tune 来优化该模型的超参数。具体来说，我们将利用提前停止（early stopping）和通过 HyperOpt 进行的贝叶斯优化来实现这一点。

提示

如果您对如何改进本教程有任何建议，请告诉我们！

要运行此示例，您需要安装以下库

$ pip install "ray[tune]" torch torchvision

设置用于调优的 PyTorch 模型#

首先，我们导入一些依赖项。我们导入一些 PyTorch 和 TorchVision 模块来帮助我们创建模型并对其进行训练。此外，我们还将导入 Ray Tune 来帮助我们优化模型。如您所见，我们使用了一个所谓的调度器（scheduler），在此例中是 ASHAScheduler，我们将在本教程后面使用它来调优模型。

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

接下来，我们定义一个简单的 PyTorch 模型供其进行训练。如果您不熟悉 PyTorch，定义模型的最简单方法是实现一个 nn.Module。这要求您使用 __init__ 设置模型，然后实现一个 forward 传递。在此示例中，我们使用了一个小型卷积神经网络，它包含一个二维卷积层、一个全连接层和一个 softmax 函数。

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)

下面，我们实现了用于训练和评估 PyTorch 模型的函数。为此，我们定义了一个 train 函数和一个 test 函数。如果您知道如何做到这一点，请跳到下一节。

训练和评估模型

# Change these values if you want the training to run quicker or slower.
EPOCH_SIZE = 512
TEST_SIZE = 256

def train_func(model, optimizer, train_loader):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # We set this just for the example to run quickly.
        if batch_idx * len(data) > EPOCH_SIZE:
            return
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()


def test_func(model, data_loader):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for batch_idx, (data, target) in enumerate(data_loader):
            # We set this just for the example to run quickly.
            if batch_idx * len(data) > TEST_SIZE:
                break
            data, target = data.to(device), target.to(device)
            outputs = model(data)
            _, predicted = torch.max(outputs.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()

    return correct / total

设置用于通过 Tune 进行训练运行的 Tuner#

下面，我们定义了一个函数，用于训练 PyTorch 模型多个 epoch。此函数将在底层通过一个单独的Ray Actor (进程)执行，因此我们需要将模型的性能传回 Tune（它在主 Python 进程中）。

为此，我们在训练函数中调用 tune.report()，它会将性能值发送回 Tune。由于该函数在单独的进程中执行，请确保该函数可以被 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)

让我们通过调用 Tuner.fit 来运行一次试验，并从均匀分布中随机采样学习率和动量。

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 返回一个 ResultGrid 对象。您可以使用它来绘制此试验的性能图。

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

注意

Tune 将自动在您机器或集群的所有可用核心/GPU 上运行并行试验。要限制并发试验的数量，请使用 ConcurrencyLimiter。

使用自适应连续减半（ASHAScheduler）进行提前停止#

让我们将提前停止集成到我们的优化过程中。我们将使用 ASHA，这是一个可扩展的算法，用于有原则的提前停止。

总的来说，ASHA 会终止不太有希望的试验，并将更多的时间和资源分配给更有希望的试验。随着我们的优化过程变得更加高效，我们可以通过调整 num_samples 参数来将搜索空间增加 5 倍。

ASHA 在 Tune 中实现为“Trial Scheduler”（试验调度器）。这些 Trial Scheduler 可以提前终止不良试验、暂停试验、克隆试验以及修改正在运行试验的超参数。有关可用调度器和库集成的更多详细信息，请参阅TrialScheduler 文档。

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}

您可以在 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)
../_images/tune-df-plot.png

您还可以使用 TensorBoard 来可视化结果。

$ tensorboard --logdir {logdir}

在 Tune 中使用搜索算法#

除了 TrialSchedulers 之外，您还可以通过使用智能搜索技术（如贝叶斯优化）来进一步优化您的超参数。为此，您可以使用 Tune 的 Search Algorithm（搜索算法）。搜索算法利用优化算法来智能地导航给定的超参数空间。

请注意，每个库都有定义搜索空间的特定方法。

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})

注意

Tune 允许您将某些搜索算法与不同的试验调度器结合使用。有关更多详细信息，请参阅此页面。

调优后评估您的模型#

您可以使用 ExperimentAnalysis 对象 来检索最佳模型，从而评估经过最佳训练的模型

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)

后续步骤#

  • 请查看 Tune 教程，了解如何将 Tune 与您喜欢的机器学习库配合使用。

  • 浏览我们的示例库，了解如何将 Tune 与 PyTorch、XGBoost、Tensorflow 等结合使用。

  • 如果您遇到问题或有任何疑问，请在我们的 GitHub 上打开一个 issue 告诉我们。

  • 要查看您的应用程序的运行情况，您可以使用Ray 仪表板。