Ray 分散式深度學習訓練與模型部署實戰

隨著資料規模和模型複雜度的提升，分散式深度學習訓練已成為解決複雜 AI 任務的關鍵。本文將以 Ray 為例，詳細介紹如何利用其生態系統進行高效的分散式訓練和模型部署。從 PyTorch 模型定義到 AIR Trainer 的使用，再到結合 Gradio 和 Ray Serve 建立預測模型服務，本文提供了一套完整的實戰。同時，我們也將比較 Dask、Spark 和 Ray 等主流分散式框架，並探討 Ray AIR 如何簡化 AI 工作流程，提升開發效率。透過本文，讀者可以深入理解 Ray 在分散式深度學習領域的應用，並掌握構建和部署高效能 AI 應用的實用技巧。

模型訓練

使用 Ray Datasets 對 CIFAR-10 進行資料預處理後，我們可以定義一個分類器來訓練資料。在這個場景中，我們將使用 PyTorch 定義一個 AIR Trainer。但是，值得注意的是，您可以輕鬆地在這個階段切換框架，使用 Keras、Hugging Face、scikit-learn 或其他由玄貓支援的庫。

我們將分三個步驟進行：定義 PyTorch 模型、指定 AIR 執行的訓練迴圈以及定義一個可以適應我們訓練資料的 AIR Trainer。首先，讓我們使用 PyTorch 定義一個簡單的卷積神經網路，具有最大池化和 ReLU 啟用函式，適合於 CIFAR-10 資料集。如果您熟悉 PyTorch，以下的神經網路定義應該很直觀。如果您不熟悉 PyTorch，只需要知道定義一個 torch.nn.Module 時，您需要提供的唯一東西是神經網路的前向傳遞定義：

import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

您已經知道，為了定義一個 AIR TorchTrainer，您需要一個訓練資料集、一個縮放因子和一個可選的執行配置。這些步驟將使您能夠使用 Ray 的 AIR 框架進行模型訓練和最佳化。下一步，我們將探討如何使用這些工具來實作高效的模型訓練和最佳化。

內容解密：

在上述程式碼中，我們定義了一個簡單的卷積神經網路，具有最大池化和 ReLU 啟用函式。這個模型的結構是為了適應 CIFAR-10 資料集而設計的。透過這個例子，我們可以看到如何使用 PyTorch 定義一個神經網路，並如何使用 Ray 的 AIR 框架進行模型訓練和最佳化。

圖表翻譯：

  graph LR
    A[資料預處理] --> B[模型定義]
    B --> C[訓練迴圈]
    C --> D[模型最佳化]
    D --> E[結果評估]

這個圖表展示了從資料預處理到結果評估的整個過程。透過這個過程，我們可以看到如何使用 Ray 的 AIR 框架進行模型訓練和最佳化。

分散式深度學習訓練迴圈

在分散式深度學習中，定義訓練迴圈是一個關鍵步驟。這涉及指定每個工作者在呼叫 fit 時應該執行的任務。為此，我們可以使用配置字典來指定執行時的屬性。

以下是使用 PyTorch 和 Ray Train 的訓練迴圈示例：

import torch
import torch.nn as nn
from ray import train
from ray.air import session, Checkpoint

def train_loop(config):
    # 載入模型
    model = train.torch.prepare_model(Net())
    
    # 定義損失函式和最佳化器
    loss_fct = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
    
    # 獲取訓練資料
    train_batches = session.get_dataset_shard("train").iter_torch_batches(
        batch_size=config["batch_size"],
    )
    
    # 訓練迴圈
    for epoch in range(config["epochs"]):
        running_loss = 0.0
        for i, data in enumerate(train_batches):
            inputs, labels = data["image"], data["label"]
            optimizer.zero_grad()
            forward_outputs = model(inputs)
            loss = loss_fct(forward_outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        
        # 報告訓練指標
        session.report({"loss": running_loss / len(train_batches)})

在這個例子中，我們定義了一個 train_loop 函式，該函式接收一個配置字典 config 作為輸入。該函式負責載入模型、定義損失函式和最佳化器、獲取訓練資料、執行訓練迴圈和報告訓練指標。

訓練迴圈的關鍵步驟

1. 載入模型：使用 train.torch.prepare_model 載入模型。
2. 定義損失函式和最佳化器：定義損失函式和最佳化器。
3. 獲取訓練資料：使用 session.get_dataset_shard 獲取訓練資料。
4. 執行訓練迴圈：執行訓練迴圈，包括前向傳播、計算損失、反向傳播和更新模型引數。
5. 報告訓練指標：使用 session.report 報告訓練指標。

Mermaid 圖表

  graph LR
    A[載入模型] --> B[定義損失函式和最佳化器]
    B --> C[獲取訓練資料]
    C --> D[執行訓練迴圈]
    D --> E[報告訓練指標]

圖表翻譯

這個圖表展示了訓練迴圈的關鍵步驟，包括載入模型、定義損失函式和最佳化器、獲取訓練資料、執行訓練迴圈和報告訓練指標。

內容解密

在這個例子中，我們使用 PyTorch 和 Ray Train 來實作分散式深度學習。訓練迴圈是分散式深度學習的核心，涉及指定每個工作者在呼叫 fit 時應該執行的任務。透過使用配置字典和 session.report 來報告訓練指標，我們可以實作分散式深度學習的訓練迴圈。

使用 Ray 進行分散式深度學習訓練

Ray 是一個高效能的分散式計算框架，提供了強大的工具和 API，用於構建和訓練大規模的深度學習模型。在這個章節中，我們將探討如何使用 Ray 進行分散式深度學習訓練。

定義訓練迴圈

首先，我們需要定義訓練迴圈，包括模型、損失函式和最佳化器。以下是使用 PyTorch 定義訓練迴圈的示例程式碼：

import torch
import torch.nn as nn
import torch.optim as optim

# 定義模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 定義損失函式和最佳化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(Net().parameters(), lr=0.01)

# 定義訓練迴圈
def train_loop(model, device, data_loader, optimizer, criterion):
    model.train()
    running_loss = 0.0
    for i, (inputs, labels) in enumerate(data_loader):
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if i % 1000 == 0:
            print(f"[{i+1:4d}] loss: {running_loss / 1000:.3f}")
            running_loss = 0.0

使用 Ray Train 進行分散式訓練

接下來，我們可以使用 Ray Train 進行分散式訓練。首先，我們需要建立一個 TorchTrainer 例項，並定義訓練迴圈和配置：

from ray.train.torch import TorchTrainer
from ray.air.callbacks.mlflow import MLflowLoggerCallback

# 定義訓練迴圈和配置
train_loop_per_worker = train_loop
train_loop_config = {"batch_size": 10, "epochs": 5}
datasets = {"train": train_dataset}

# 建立 TorchTrainer 例項
trainer = TorchTrainer(
    train_loop_per_worker=train_loop,
    train_loop_config=train_loop_config,
    datasets=datasets,
    scaling_config=ScalingConfig(num_workers=2),
    run_config=RunConfig(callbacks=[MLflowLoggerCallback(experiment_name="torch_trainer")])
)

# 進行訓練
result = trainer.fit()

使用 MLflow 進行日誌記錄

在上面的例子中，我們使用了 MLflowLoggerCallback 進行日誌記錄。這允許我們將訓練結果記錄到 MLflow 中，以便於追蹤和分析訓練過程。

Ray Train 和 RLlib 生態系統

Ray Train 和 RLlib 提供了強大的工具和 API，用於構建和訓練大規模的深度學習模型。以下是 Ray Train 和 RLlib 生態系統的總結：

圖表翻譯：

  graph LR
    A[Ray Train] --> B[RLlib]
    A --> C[MLflow]
    A --> D[Weights & Biases]
    A --> E[CometML]
    B --> C
    B --> D
    B --> E
    C --> D
    C --> E
    D --> E

在這個圖表中，我們可以看到 Ray Train 和 RLlib 是兩個核心框架，提供了強大的工具和 API，用於構建和訓練大規模的深度學習模型。MLflow、Weights & Biases 和 CometML 是三個日誌記錄和管理框架，提供了強大的工具和 API，用於記錄和分析訓練過程。

人工智慧模型部署與服務

人工智慧模型的部署與服務是將訓練好的模型應用於實際場景的重要步驟。為了簡化這個過程，開發者們建立了各種工具和框架，以便更容易地部署和管理模型。其中，Gradio是一個流行的工具，允許開發者輕鬆地建立圖形使用者介面，以展示和測試他們的機器學習模型。

Gradio的簡單使用

Gradio提供了一個簡單的方式來建立圖形使用者介面，讓使用者可以與機器學習模型進行互動。開發者可以使用Gradio的API來定義模型的輸入和輸出，然後Gradio會自動建立一個圖形使用者介面。例如，以下程式碼展示瞭如何使用Gradio建立一個簡單的圖形使用者介面：

import gradio as gr
import numpy as np

def predict(payload):
    # 將輸入轉換為NumPy陣列
    payload = np.array(payload, dtype=np.float32)
    # 將陣列重塑為模型的輸入形狀
    array = payload.reshape((1, 3, 32, 32))
    # 使用模型進行預測
    return np.argmax(predictor.predict(array))

# 建立Gradio介面
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Label(label="預測結果")
)

# 啟動Gradio介面
demo.launch()

這個程式碼建立了一個Gradio介面，允許使用者上傳圖片，並使用機器學習模型進行預測。Gradio會自動建立一個圖形使用者介面，讓使用者可以與模型進行互動。

Ray Serve的簡單使用

Ray Serve是一個根據Ray的模型服務框架，允許開發者輕鬆地部署和管理機器學習模型。Ray Serve提供了一個簡單的方式來建立模型服務，讓開發者可以將模型部署到雲端或本地環境。以下程式碼展示瞭如何使用Ray Serve建立一個簡單的模型服務：

import ray
from ray.serve import serve

# 建立Ray Serve應用
@serve.deployment(route_prefix="/predict")
def predict(request):
    # 將輸入轉換為NumPy陣列
    payload = np.array(request.json["payload"], dtype=np.float32)
    # 將陣列重塑為模型的輸入形狀
    array = payload.reshape((1, 3, 32, 32))
    # 使用模型進行預測
    return np.argmax(predictor.predict(array))

# 啟動Ray Serve應用
ray.init()
serve.start(detached=True)

# 建立Gradio介面
demo = gr.Interface(
    fn=lambda x: requests.post("http://localhost:8000/predict", json={"payload": x}).json(),
    inputs=gr.Image(type="numpy"),
    outputs=gr.Label(label="預測結果")
)

# 啟動Gradio介面
demo.launch()

這個程式碼建立了一個Ray Serve應用，允許使用者傳送請求到模型服務，並使用機器學習模型進行預測。Ray Serve會自動建立一個模型服務，讓開發者可以將模型部署到雲端或本地環境。

結合Gradio和Ray Serve

Gradio和Ray Serve可以結合使用，允許開發者輕鬆地建立圖形使用者介面和模型服務。以下程式碼展示瞭如何結合Gradio和Ray Serve：

import gradio as gr
import numpy as np
import ray
from ray.serve import serve

# 建立Ray Serve應用
@serve.deployment(route_prefix="/predict")
def predict(request):
    # 將輸入轉換為NumPy陣列
    payload = np.array(request.json["payload"], dtype=np.float32)
    # 將陣列重塑為模型的輸入形狀
    array = payload.reshape((1, 3, 32, 32))
    # 使用模型進行預測
    return np.argmax(predictor.predict(array))

# 啟動Ray Serve應用
ray.init()
serve.start(detached=True)

# 建立Gradio介面
demo = gr.Interface(
    fn=lambda x: requests.post("http://localhost:8000/predict", json={"payload": x}).json(),
    inputs=gr.Image(type="numpy"),
    outputs=gr.Label(label="預測結果")
)

# 啟動Gradio介面
demo.launch()

這個程式碼結合了Gradio和Ray Serve，允許開發者輕鬆地建立圖形使用者介面和模型服務。Gradio會自動建立一個圖形使用者介面，讓使用者可以與模型進行互動，而Ray Serve會自動建立一個模型服務，讓開發者可以將模型部署到雲端或本地環境。

使用 Ray Serve 和 Gradio 建立預測模型服務

Ray Serve 是一個強大的服務框架，允許您輕鬆地建立和部署機器學習模型。Gradio 是一個開源的機器學習模型服務框架，允許您建立和部署機器學習模型。在這個範例中，我們將使用 Ray Serve 和 Gradio 建立一個預測模型服務。

建立 Gradio 介面

首先，我們需要建立一個 Gradio 介面，該介面將接受圖片輸入並輸出預測結果。

import gradio as gr

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=10)
)

在這個範例中，predict 是一個預測函式，該函式接受圖片輸入並輸出預測結果。gr.Image() 是 Gradio 的圖片輸入元件，gr.Label() 是 Gradio 的標籤輸出元件。

建立 Ray Serve 應用

接下來，我們需要建立一個 Ray Serve 應用，該應用將使用 Gradio 介面。

from ray import serve

app = serve.start(detached=True)

在這個範例中，serve.start() 是 Ray Serve 的啟動函式，該函式啟動 Ray Serve 服務。

繫結 Gradio 介面到 Ray Serve 應用

現在，我們需要繫結 Gradio 介面到 Ray Serve 應用。

from ray.serve import RayServe

ray_serve = RayServe()
ray_serve.bind(demo)

在這個範例中，RayServe() 是 Ray Serve 的啟動函式，該函式啟動 Ray Serve 服務。bind() 是 Ray Serve 的繫結函式，該函式繫結 Gradio 介面到 Ray Serve 應用。

執行 Ray Serve 應用

最後，我們可以執行 Ray Serve 應用。

serve run gradio_demo:app

在這個範例中，serve run 是 Ray Serve 的執行命令，該命令啟動 Ray Serve 服務。gradio_demo:app 是 Ray Serve 的應用名稱，該名稱對應於我們建立的 Gradio 介面。

測試 Ray Serve 應用

現在，我們可以測試 Ray Serve 應用。首先，我們需要上傳圖片到 Ray Serve 應用。

curl -X POST \
  http://localhost:8000/predict \
  -H 'Content-Type: application/json' \
  -d '{"image": "image.jpg"}'

在這個範例中，curl 是一個命令列工具，該工具用於傳送 HTTP 請求。http://localhost:8000/predict 是 Ray Serve 的預測 URL，該 URL 對應於我們建立的 Gradio 介面。-H 'Content-Type: application/json' 是 HTTP 請求的標頭，該標頭指定請求的內容型別為 JSON。-d '{"image": "image.jpg"}' 是 HTTP 請求的請求體，該請求體包含圖片檔案的路徑。

結果

如果一切順利，Ray Serve 應用將輸出預測結果。

{
  "label": "dog"
}

在這個範例中，label 是預測結果，該結果對應於圖片的類別。

分散式Python框架

在探討分散式計算框架時，若考慮到完全支援Python且不鎖定任何雲端服務，目前的主流框架有Dask、Spark和Ray。雖然這些框架在技術上和特定情境下的效能表現可能有所不同，但最好的比較方法是根據您想要在其上執行的工作負載型別。

工作負載型別比較

下表比較了Dask、Spark和Ray在不同工作負載型別上的支援情況：

框架選擇

選擇框架時，需要考慮您的具體需求和工作負載特點。例如，如果您主要處理結構化資料，Dask和Spark可能是更好的選擇。然而，如果您需要一個更通用的平臺，可以支援從機器學習到實時資料處理的多種工作負載，Ray可能是更好的選擇。

Ray的生態系統

Ray的生態系統提供了豐富的工具和整合，包括Ray Tune、Ray AIR和Ray Datasets等。這些工具可以幫助您更高效地開發和部署機器學習模型和其他分散式應用。

內容解密：

上述內容介紹了分散式Python框架的比較，包括Dask、Spark和Ray。每個框架都有其優勢和劣勢，選擇框架時需要考慮工作負載型別和具體需求。Ray的生態系統提供了豐富的工具和整合，可以幫助開發和部署機器學習模型和其他分散式應用。

  flowchart TD
    A[工作負載型別] --> B[選擇框架]
    B --> C[Dask]
    B --> D[Spark]
    B --> E[Ray]
    C --> F[結構化資料處理]
    D --> G[結構化資料處理]
    E --> H[機器學習和實時資料處理]

圖表翻譯：

此圖表示工作負載型別與框架選擇之間的關係。根據工作負載型別的不同，可以選擇適合的框架。Dask和Spark更適合結構化資料處理，而Ray則提供了更通用的平臺，可以支援多種工作負載，包括機器學習和實時資料處理。

人工智慧工作流程與 Ray AIR

人工智慧（AI）工作流程涉及多個階段，包括資料準備、模型訓練、模型評估、模型部署等。Ray AIR是一個根據Ray的AI計算框架，旨在簡化AI工作流程的管理和執行。

Ray AIR 與其他框架的比較

Ray AIR 與其他框架相比，有以下幾個優點：

• 框架無關性：Ray AIR 不依賴於特定的深度學習框架，可以與多個框架（如 PyTorch、TensorFlow 等）無縫整合。
• 分散式計算：Ray AIR 支援分散式計算，可以將計算任務分配到多個機器上，從而提高計算效率。
• 簡單易用：Ray AIR 提供了一個簡單易用的 API，讓使用者可以輕鬆地管理和執行 AI 工作流程。

Ray AIR 的生態系統

Ray AIR 的生態系統包括多個元件，例如：

• ML 跟蹤和觀察：MLflow、Weights & Biases、Arize 等。
• 訓練框架：PyTorch、TensorFlow、Lightning、JAX 等。
• ML 特徵儲存：Feast、Tecton 等。

從技術架構視角來看，Ray 生態系統的發展，特別是 Ray AIR 的出現，為機器學習工作流程的建構和管理提供了一個強大的解決方案。透過整合 PyTorch、TensorFlow 等主流深度學習框架，以及 MLflow、Weights & Biases 等日誌記錄和管理工具，Ray AIR 簡化了從資料預處理、模型訓練到模型部署的整個流程。分析其核心元件，如 Ray Datasets、Ray Tune 和 Ray Serve，可以發現 Ray 生態系統致力於解決分散式機器學習中的關鍵挑戰，例如資料分片、超引數調整和模型服務。然而，目前 Ray 生態系統的學習曲線相對較陡峭，需要開發者投入一定的時間和精力去學習和掌握。對於重視開發效率的團隊，需要權衡學習成本和長期收益。玄貓認為，Ray 生態系統代表了未來機器學習平臺發展的重要方向，值得投入資源深入研究和應用，尤其對於需要處理大規模資料和複雜模型的團隊而言，更能發揮其優勢。隨著社群的持續發展和工具的日漸成熟，我們預見 Ray 的應用門檻將逐步降低，其影響力也將持續擴大。