PyTorch & 分布式框架 Ray :保姆級入門教程

使用 PyTorch 和 Ray 進行分布式機器學習入門，Ray 是一個流行的分布式Python 框架，可以與 PyTorch 搭配，快速擴展機器學習應用。

來源：官方博客

翻譯：PyTorch 開發者社區（微信公眾號）

今天的機器學習需要分布式計算。無論是訓練網絡、調整超參數、服務模型還是處理數據，機器學習都是計算密集型的，如果沒有訪問集群，速度會非常慢。

Ray 是一個流行的分布式 Python 框架，它可以與 PyTorch 配對，以快速擴展機器學習應用。

本篇文章介紹 Ray 生態系統的各種元素，以及如何與 PyTorch 搭配使用!

Ray 是一個並行和分布式 Python 的開源庫。

從高層次上看，Ray 生態系統由三部分組成：核心 Ray 系統、用於機器學習的可擴展庫（包括原生庫和第三方庫），以及用於在任何集群或雲提供商上啟動集群的工具。

Ray 可以用來在多個核心或機器上擴展 Python 應用。它有幾個主要的優點，包括：

簡單性：你可以擴展你的 Python 應用，而不需要重寫，同樣的代碼可以在一臺機器或多臺機器上運行。

穩健性：應用程式可以優雅地處理機器故障和進程搶佔。

性能：任務以毫秒級的延遲運行，可擴展到數萬個內核，並以最小的序列化開銷處理數值數據。

由於Ray是一個通用框架，社區在它的基礎上建立了許多庫和框架來完成不同的任務。

這些庫和框架絕大多數都支持 PyTorch，只需對代碼進行最小程度的修改，並能相互無縫集成。以下是生態系統中眾多庫中的一部分。

圖註：在 p3dn.24xlarge 實例上，PyTorch 的 DataParallel 與 Ray 的比較（Ray 在下面使用 PyTorch 的分布式 DataParallel）。

RaySGD是一個為數據並行訓練提供分布式訓練包裝的庫。例如，RaySGD TorchTrainer（https://docs.ray.io/en/master/raysgd/raysgd_pytorch.html） 是一個圍繞 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 torchfrom torch.utils.data import DataLoaderfrom torchvision.datasets import CIFAR10import torchvision.transforms as transforms
import rayfrom ray.util.sgd.torch import TorchTrainer# https: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 npimport torchimport torch.optim as optim
from ray import tunefrom ray.tune.examples.mnist_pytorch import get_data_loaders, train, testimport rayimport sys
if len(sys.argv) > 1:    ray.init(redis_address=sys.argv[1])
import torch.nn as nnimport 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 不僅可以用來單獨為模型服務，還可以用來擴展其他服務工具，比如 FastAPI。

Ray Serve 是一個易於使用的可擴展模型服務的庫。該庫的一些優點是：
能夠使用一個工具包來服務從深度學習模型（PyTorch，TensorFlow等）到 scikit-learn 模型，再到任意的 Python 業務邏輯。
可擴展到許多機器，無論是在你的數據中心還是在雲端。
與許多其他庫兼容，如 Ray Tune 和 FastAPI。
如果您想了解如何將 Ray Serve 和 Ray Tune 一起整合到您的 PyTorch 工作流中，您應該查看文檔以獲取完整的代碼示例。
圖註：RLlib提供了幾乎所有訓練方面的定製方法，包括神經網絡模型、動作分布、策略定義、環境和樣本收集過程。

RLlib提供了幾乎所有訓練方面的定製方法，包括神經網絡模型、動作分布、策略定義、環境和樣本收集過程。
RLlib 是一個用於強化學習的庫，它既提供了高擴展性，又為各種應用提供了統一的 API，優勢包括：
圖註：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 加入項目的社區。