开始使用 Ray Tune#

本教程将引导您完成设置 Tune 实验的过程。首先,我们采用一个 PyTorch 模型,并向您展示如何利用 Ray Tune 来优化该模型的超参数。具体来说,我们将通过 HyperOpt 利用早期停止和贝叶斯优化来实现这一点。

小技巧

如果您有关于如何改进本教程的建议,请 告诉我们!

要运行此示例,您需要安装以下内容:

$ pip install "ray[tune]" torch torchvision

设置一个 Pytorch 模型以进行调优#

首先,让我们先导入一些依赖项。我们导入一些 PyTorch 和 TorchVision 模块来帮助我们创建模型并训练它。此外,我们将导入 Ray Tune 来帮助我们优化模型。如你所见,我们使用了一个所谓的调度器,在本例中是 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 train, tune
from ray.tune.schedulers import ASHAScheduler

接下来,让我们定义一个简单的 PyTorch 模型,我们将对其进行训练。如果你不熟悉 PyTorch,定义模型的最简单方法是实现一个 nn.Module。这需要你使用 __init__ 设置模型,然后实现一个 forward 传递。在这个例子中,我们使用了一个由一个 2D 卷积层、一个全连接层和一个 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 模型多个周期。这个函数将在底层的一个单独的 Ray Actor (进程) 上执行,因此我们需要将模型的性能反馈给 Tune(它在主 Python 进程上)。

为此,我们在训练函数中调用 train.report(),这将性能值发送回 Tune。由于该函数在单独的进程中执行,请确保该函数 可由 Ray 序列化

import os
import tempfile

from ray.train 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
            train.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 上运行并行试验。要限制并发试验的数量,请使用 并发限制器

使用自适应连续减半的早期停止 (ASHAScheduler)#

让我们将早停机制整合到我们的优化过程中。让我们使用 ASHA,这是一种用于 principled early stopping 的可扩展算法。

在高层次上,ASHA 终止了那些不太有希望的试验,并将更多的时间和资源分配给更有希望的试验。随着我们的优化过程变得更加高效,我们可以通过调整参数 num_samples将搜索空间增加5倍

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

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 笔记本中运行以下代码来可视化试验进度。

# 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 搜索算法 。搜索算法利用优化算法来智能地导航给定的超参数空间。

请注意,每个库都有其特定的定义搜索空间的方式。

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 允许你结合不同的试验调度器使用一些搜索算法。详情请参阅 此页面

调整后评估您的模型#

您可以使用 实验分析对象 来评估最佳训练模型以检索最佳模型:

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 仪表盘