PyTorch & 分布式框架 Ray ：保姆级入门教程

来源：官方博客
翻译：PyTorch 开发者社区（微信公众号）

今天的机器学习需要分布式计算。无论是训练网络、调整超参数、服务模型还是处理数据，机器学习都是计算密集型的，如果没有访问集群，速度会非常慢。

Ray 是一个流行的分布式 Python 框架，它可以与 PyTorch 配对，以快速扩展机器学习应用。

本篇文章介绍 Ray 生态系统的各种元素，以及如何与 PyTorch 搭配使用!

Ray 是什么？

Ray 是一个并行和分布式 Python 的开源库。

从高层次上看，Ray 生态系统由三部分组成：核心 Ray 系统、用于机器学习的可扩展库（包括原生库和第三方库），以及用于在任何集群或云提供商上启动集群的工具。

Ray 的核心系统

Ray 可以用来在多个核心或机器上扩展 Python 应用。它有几个主要的优点，包括：

简单性：你可以扩展你的 Python 应用，而不需要重写，同样的代码可以在一台机器或多台机器上运行。

稳健性：应用程序可以优雅地处理机器故障和进程抢占。

性能：任务以毫秒级的延迟运行，可扩展到数万个内核，并以最小的序列化开销处理数值数据。

Ray 的 library 生态

由于Ray是一个通用框架，社区在它的基础上建立了许多库和框架来完成不同的任务。

这些库和框架绝大多数都支持 PyTorch，只需对代码进行最小程度的修改，并能相互无缝集成。以下是生态系统中众多库中的一部分。

RaySGD

图注：在 p3dn.24xlarge 实例上，PyTorch 的 DataParallel 与 Ray 的比较（Ray 在下面使用
PyTorch 的分布式 DataParallel）。

RaySGD 是一个为数据并行训练提供分布式训练包装的库。例如，RaySGD TorchTrainer（https://docs.ray.io/en/master...） 是一个围绕 torch.distributed.launch 的包装器。它提供了一个 Python API，可以轻松地将分布式训练纳入到一个更大的 Python 应用程序中，而不是需要将你的训练代码包装在 bash 脚本中。

该库的其他一些优点是：

易用性：您可以扩展 PyTorch 的原生 DistributedDataParallel，而无需监控单个节点。

可扩展性：您可以对 PyTorch 的原生分布式数据并行进行扩展，而无需监控单个节点。您可以向上和向下扩展。从单个 CPU 开始。只需更改两行代码，即可扩展到多节点、多 CPU 或多 GPU 集群。

加速训练：NVIDIA Apex 内置了对混合精度训练的支持。

容错功能：支持在云机被抢占时自动恢复。

兼容性：支持与其他库无缝集成，如：NVIDIA Apex。可与其他库无缝集成，如Ray Tune 和 Ray Serve。

你可以通过安装 Ray（pip install -U ray torch）并运行下面的代码，来开始使用 TorchTrainer：

import torch
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
import torchvision.transforms as transforms
​
import ray
from ray.util.sgd.torch import TorchTrainer
# https://github.com/kuangliu/pytorch-cifar/blob/master/models/resnet.py
from ray.util.sgd.torch.resnet import ResNet18
​
​
def cifar_creator(config):
    """Returns dataloaders to be used in `train` and `validate`."""
    tfms = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465),
                             (0.2023, 0.1994, 0.2010)),
    ])  # meanstd transformation
    train_loader = DataLoader(
        CIFAR10(root="~/data", download=True, transform=tfms), batch_size=config["batch"])
    validation_loader = DataLoader(
        CIFAR10(root="~/data", download=True, transform=tfms), batch_size=config["batch"])
    return train_loader, validation_loader
​
​
def optimizer_creator(model, config):
    """Returns an optimizer (or multiple)"""
    return torch.optim.SGD(model.parameters(), lr=config["lr"])
​
ray.init()
​
trainer = TorchTrainer(
    model_creator=ResNet18,  # A function that returns a nn.Module
    data_creator=cifar_creator,  # A function that returns dataloaders
    optimizer_creator=optimizer_creator,  # A function that returns an optimizer
    loss_creator=torch.nn.CrossEntropyLoss,  # A loss function
    config={"lr": 0.01, "batch": 64},  # parameters
    num_workers=2,  # amount of parallelism
    use_gpu=torch.cuda.is_available(),
    use_tqdm=True)
​
stats = trainer.train()
print(trainer.validate())
​
torch.save(trainer.state_dict(), "checkpoint.pt")
trainer.shutdown()
print("success!")

注释：该脚本将下载 CIFAR10 并使用 ResNet18 模型进行图像分类。只需更改一个参数（num_workers=N），就可以利用多个 GPU。

Ray Tune

图注：Ray Tune 实现的优化算法，如基于群体的训练（如上图所示），可以与 PyTorch 一起使用，以获得更高性能的模型。

Ray Tune 是一个 Python 库，用于任何规模的实验执行和超参数调整，该库的一些优点是：

• 能够在不到 10 行代码中启动多节点分布式超参数扫描。
• 支持包括 PyTorch 在内的所有主流机器学习框架。
• 对 GPU 的一流支持。
• 自动管理检查点并将日志记录到 TensorBoard。
• 访问最先进的算法，如基于群体的训练（PBT）、BayesOptSearch、HyperBand/ASHA。

你可以通过安装 Ray（pip install ray torch torchvision）并运行下面的代码，来开始使用 Ray Tune。

import numpy as np
import torch
import torch.optim as optim
​
from ray import tune
from ray.tune.examples.mnist_pytorch import get_data_loaders, train, test
import ray
import sys
​
if len(sys.argv) > 1:
    ray.init(redis_address=sys.argv[1])
​
import torch.nn as nn
import torch.nn.functional as F
​
class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        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)
​
​
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.state_dict(), "./model.pth")
​
from ray.tune.schedulers import ASHAScheduler
​
search_space = {
    "lr": tune.choice([0.001, 0.01, 0.1]),
    "momentum": tune.uniform(0.1, 0.9)
}
​
analysis = tune.run(
    train_mnist,
    num_samples=30,
    scheduler=ASHAScheduler(metric="mean_accuracy", mode="max", grace_period=1),
    config=search_space)

注释：该脚本向你展示了如何利用最先进的早期停止算法 AHSA，它可以终止那些不太有前途的试验，并将更多的时间和资源分配给更有前途的试验。

Ray Serve

图注：Ray Serve 不仅可以用来单独为模型服务，还可以用来扩展其他服务工具，比如 FastAPI。

Ray Serve 是一个易于使用的可扩展模型服务的库。该库的一些优点是：

能够使用一个工具包来服务从深度学习模型（PyTorch，TensorFlow等）到 scikit-learn 模型，再到任意的 Python 业务逻辑。
可扩展到许多机器，无论是在你的数据中心还是在云端。
与许多其他库兼容，如 Ray Tune 和 FastAPI。

如果您想了解如何将 Ray Serve 和 Ray Tune 一起整合到您的 PyTorch 工作流中，您应该查看文档以获取完整的代码示例。

RLlib

图注：RLlib 提供了几乎所有训练方面的定制方法，包括神经网络模型、动作分布、策略定义、环境和样本收集过程。

RLlib 提供了几乎所有训练方面的定制方法，包括神经网络模型、动作分布、策略定义、环境和样本收集过程。

RLlib 是一个用于强化学习的库，它既提供了高扩展性，又为各种应用提供了统一的 API，优势包括：

• 原生支持 PyTorch、TensorFlow Eager 和 TensorFlow（1.x和2.x）。
• 支持无模型、基于模型、进化、规划和多代理算法。
• 通过简单的配置标志和自动封装器支持复杂的模型类型，如注意力网和LSTM 堆栈。
• 与 Ray Tune 等其他库的兼容性。

Cluster Launcher

图注：Ray Cluster Launcher简化了在任何集群或云提供商上启动和扩展的过程。

一旦您在笔记本电脑上开发了一个应用程序，并希望将其扩展到云端（也许有更多的数据或更多的 GPU），接下来的步骤并不总是很清楚。这个过程要么让基础设施团队为你设置，要么通过以下步骤:

1. 选择一个云提供商（AWS、GCP 或 Azure）。
2. 导航管理控制台，设置实例类型、安全策略、节点、实例限制等。
3. 弄清楚如何在集群上分发你的 Python 脚本。

一个更简单的方法是使用 Ray Cluster Launcher 在任何集群或云提供商上启动和扩展机器。Cluster Launcher 允许你自动缩放、同步文件、提交脚本、端口转发等。这意味着您可以在 Kubernetes、AWS、GCP、Azure 或私有集群上运行您的 Ray 集群，而无需了解集群管理的低级细节。

## 总结

图注：Ray 为蚂蚁金服集团的 Fusion Engine 提供了一个分布式计算基础。

本文包含了 Ray 在 PyTorch 生态系统中的一些好处。Ray 被广泛用于各种应用，从蚂蚁金服集团使用 Ray 支持其金融业务，到 LinkedIn 在 Yarn 上运行Ray，再到 Pathmind 使用 Ray 将强化学习与模拟软件连接起来，等等。

如果你有任何关于 Ray 的问题或想法，或者想了解更多关于并行和分布式Python 的信息，请通过 Discourse、Slack 或 GitHub 加入项目的社区。