分布式訓練 DLC 快速入門 - Platform For AI

DLC可以快捷地建立分布式或單機訓練任務。其底層基於Kubernetes，省去您手動購買機器並配置運行環境，無需改變使用習慣即可快速使用。本文以 MNIST 手寫體識別為例，介紹如何使用DLC進行單機單卡訓練，或多機多卡的分布式訓練。

說明

MNIST手寫體識別是深度學習最經典的入門任務之一，任務目標是通過構建機器學習模型，來識別10個手寫數字（0～9）。

前提條件

使用主帳號開通PAI並建立工作空間。登入PAI控制台，左上方選擇開通地區，然後一鍵授權和開通產品。

計費說明

本文案例將使用公用資源建立DLC任務，計費方式為隨用隨付，詳細計費規則請參見分布式訓練（DLC）計費說明。

單機單卡訓練

建立資料集

資料集用於儲存模型訓練的代碼、資料、以及訓練結果。本文以Object Storage Service類型資料集為例進行說明。

在PAI控制台左側功能表列單擊資料集 > 自訂資料集 > 建立資料集。
配置資料集參數。關鍵參數配置如下，其他參數預設即可。
- 名稱：如：dataset_mnist
- 儲存類型：Object Storage Service
- OSS路徑：單擊表徵圖，選擇Bucket並建立目錄如：dlc_mnist。
  如果您尚未開通OSS，或在當前地區下沒有可選的Bucket，可參考如下步驟開通OSS，並建立Bucket：
  （可選）開通OSS，並建立Bucket
  1. 開通OSS服務。
  2. 登入OSS管理主控台，單擊建立Bucket，填寫Bucket名稱，地區選擇與當前PAI相同的地區，其他參數預設即可，然後單擊完成建立。
單擊確定建立資料集。

上傳訓練代碼和資料。

下載代碼。本文已經為您準備好了訓練代碼，單擊mnist_train.py下載。為減少您的操作，代碼運行時會自動將訓練資料下載到資料集的dataSet目錄中。

您在後續實際業務使用時，可以預先把代碼和訓練資料上傳到PAI的資料集中。

單機單卡訓練程式碼範例 mnist_train.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.utils.tensorboard import SummaryWriter

# 超參數
batch_size = 64  # 每次訓練的資料量
learning_rate = 0.01  # 學習率
num_epochs = 20  # 訓練輪次

# 檢查是否有可用的 GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 資料預先處理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_dataset = datasets.MNIST(root='/mnt/data/dataSet', train=True, download=True, transform=transform)
val_dataset = datasets.MNIST(root='/mnt/data/dataSet', train=False, download=False, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)


# 定義簡單的神經網路
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # 第一層卷積：輸入通道1（灰階映像），輸出通道10，卷積核5x5
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        # 第二層卷積：輸入通道10，輸出通道20，卷積核3x3
        self.conv2 = nn.Conv2d(10, 20, kernel_size=3)
        # 全串連層：輸入為20*5*5（卷積+池化後的特徵圖尺寸），輸出128
        self.fc1 = nn.Linear(20 * 5 * 5, 128)
        # 輸出層：128 -> 10（對應10個數字類別）
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        # 輸入x形狀: [batch, 1, 28, 28]
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)  # [batch, 10, 12, 12]
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)  # [batch, 20, 5, 5]
        x = x.view(-1, 20 * 5 * 5)  # 展平為[batch, 500]
        x = F.relu(self.fc1(x))      # [batch, 128]
        x = self.fc2(x)              # [batch, 10]
        return x


# 執行個體化模型，並將其移動到 GPU 上（如果可用）
model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)

# 建立 TensorBoard 的 SummaryWriter，可用於可視化的查看模型訓練過程
writer = SummaryWriter('/mnt/data/output/runs/mnist_experiment')

# 用於儲存最高準確率的模型的變數
best_val_accuracy = 0.0

# 訓練模型並記錄損失和準確率
for epoch in range(num_epochs):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)  # 將資料和目標移動到 GPU

        # 清零梯度
        optimizer.zero_grad()
        # 前向傳播
        output = model(data)
        # 計算損失
        loss = criterion(output, target)
        # 反向傳播
        loss.backward()
        # 更新參數
        optimizer.step()

        # 記錄訓練損失到 TensorBoard
        if batch_idx % 100 == 0:  # 每 100 個批次記錄一次
            writer.add_scalar('Loss/train', loss.item(), epoch * len(train_loader) + batch_idx)
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')

    # 驗證模型並記錄驗證損失和準確率
    model.eval()
    val_loss = 0
    correct = 0
    with torch.no_grad():  # 不計算梯度
        for data, target in val_loader:
            data, target = data.to(device), target.to(device)  # 將資料和目標移動到 GPU
            output = model(data)
            val_loss += criterion(output, target).item()  # 累加驗證損失
            pred = output.argmax(dim=1, keepdim=True)  # 擷取預測標籤
            correct += pred.eq(target.view_as(pred)).sum().item()  # 累加正確預測的數量

    val_loss /= len(val_loader)  # 計算平均驗證損失
    val_accuracy = 100. * correct / len(val_loader.dataset)  # 計算驗證準確率
    print(f'Validation Loss: {val_loss:.4f}, Accuracy: {correct}/{len(val_loader.dataset)} ({val_accuracy:.0f}%)')

    # 記錄驗證損失和準確率到 TensorBoard
    writer.add_scalar('Loss/validation', val_loss, epoch)
    writer.add_scalar('Accuracy/validation', val_accuracy, epoch)

    # 儲存驗證準確率最高的模型
    if val_accuracy > best_val_accuracy:
        best_val_accuracy = val_accuracy
        torch.save(model.state_dict(), '/mnt/data/output/best_model.pth')
        print(f'Model saved with accuracy: {best_val_accuracy:.2f}%')

# 關閉 SummaryWriter
writer.close()
print('Training complete. writer.close()')

上傳代碼。在資料集詳情頁，單擊查看資料跳轉至OSS控制台。然後單擊上傳檔案 > 掃描檔案 > 上傳檔案，將訓練代碼上傳至OSS中。

建立DLC任務

在PAI控制台左側功能表列單擊分布式訓練DLC > 建立任務。

配置DLC任務參數。關鍵參數配置如下，其他參數預設即可。全量參數請參見建立訓練任務。

鏡像配置：選擇鏡像地址，然後根據您所在地區填寫對應鏡像地址。

地區	對應鏡像地址
北京	dsw-registry-vpc.cn-beijing.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04
上海	dsw-registry-vpc.cn-shanghai.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04
杭州	dsw-registry-vpc.cn-hangzhou.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04
其他	查詢地區ID，並替換鏡像地址中的<地區ID>擷取完整連結： dsw-registry-vpc.<地區ID>.cr.aliyuncs.com/pai/modelscope:1.28.0-pytorch2.3.1tensorflow2.16.1-gpu-py311-cu121-ubuntu22.04

該鏡像已在互動式建模 DSW 快速入門中驗證沒有環境問題。使用PAI建模時，通常先在DSW中驗證環境、開發代碼，然後再使用DLC訓練。

資料集掛載：選擇自訂資料集，選擇上一步中建立的資料集。掛載路徑預設/mnt/data。
啟動命令：python /mnt/data/mnist_train.py
該啟動命令與在DSW或本地運行時相同。但由於mnist_train.py 現已掛載至 /mnt/data/，因此僅需要修改代碼的路徑為/mnt/data/mnist_train.py。
資源來源：選擇公用資源，資源規格選擇ecs.gn7i-c8g1.2xlarge即可。
如果該規格執行個體庫存不足，您也可以選擇其他GPU執行個體。

單擊確定建立任務，任務大約需要執行15分鐘。執行過程中可以單擊日誌查看訓練過程。

執行完成後，會在掛載資料集的output路徑下輸出最佳的模型檢查點，以及TensorBoard日誌。

（可選）查看Tensorboard

您可以藉助視覺化檢視TensorBoard查看loss曲線，瞭解訓練的具體情況。

重要

DLC任務如果想使用TensorBoard，必須配置資料集。

單擊DLC任務詳情頁上方的Tensorboard > 建立Tensorboard。
配置類型選擇按任務，在Summary目錄處填寫訓練代碼中Summary儲存的路徑：/mnt/data/output/runs/，單擊確定啟動。
對應程式碼片段：writer = SummaryWriter('/mnt/data/output/runs/mnist_experiment')
單擊查看Tensorboard查看train_loss曲線（反映訓練集損失）與 validation_loss曲線（反映驗證集損失）。
（可選）根據loss映像，調整超參數，提升模型效果
您可以根據損失值的變化趨勢，初步判斷當前模型的訓練效果：
- 在結束訓練前 train_loss 與 validation_loss 仍有下降趨勢（欠擬合）
  您可以增加 num_epochs（訓練輪次，與訓練深度正相關），或適當增大 learning_rate 後再進行訓練，加大模型的對訓練資料的擬合程度；
- 在結束訓練前 train_loss 持續下降，validation_loss 開始變大（過擬合）
  您可以減少 num_epochs，或適當減小 learning_rate 後再進行訓練，防止模型過度訓練；
- 在結束訓練前 train_loss 與 validation_loss 均處於平穩狀態（良好擬合）
  模型處於該狀態時，您可以進行後續步驟。

部署訓練後的模型

詳情請參見使用EAS將模型部署為線上服務。