Linux環境下PyTorch如何進行分布式訓練

在Linux環境下，使用PyTorch進行分布式訓練通常涉及以下幾個步驟：

1. 環境準備：

   • 確保所有節點上都安裝了相同版本的PyTorch和依賴庫。
   • 確保所有節點可以通過SSH無密碼登錄。

2. 啟動分布式訓練： PyTorch提供了torch.distributed.launch工具來啟動分布式訓練。你需要指定一些參數，如總的GPU數量、每個節點的GPU數量、節點的IP地址和端口號等。

   python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_PER_NODE \
       --nnodes=NUM_NODES --node_rank=NODE_RANK --master_addr=MASTER_IP \
       --master_port=MASTER_PORT YOUR_TRAINING_SCRIPT.py
   

   參數說明：

   • --nproc_per_node：每個節點上使用的GPU數量。
   • --nnodes：總的節點數量。
   • --node_rank：當前節點的排名（從0開始）。
   • --master_addr：主節點的IP地址。
   • --master_port：主節點的端口號。
   • YOUR_TRAINING_SCRIPT.py：你的訓練腳本。

3. 修改訓練腳本： 在你的訓練腳本中，需要初始化分布式環境。通常在腳本的最開始添加以下代碼：

   import torch.distributed as dist
   
   dist.init_process_group(backend='nccl',  # 或者 'gloo'
                           init_method='tcp://MASTER_IP:MASTER_PORT',
                           world_size=WORLD_SIZE,  # 總的進程數
                           rank=RANK)  # 當前進程的排名
   

   參數說明：

   • backend：分布式后端，常用的有nccl（用于GPU）和gloo（用于CPU和GPU）。
   • init_method：初始化分布式環境的地址。
   • world_size：總的進程數。
   • rank：當前進程的排名。

4. 數據并行： 在訓練腳本中，使用torch.nn.parallel.DistributedDataParallel來包裝你的模型，以實現數據并行。

   model = YourModel().to(rank)
   ddp_model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank])
   

5. 數據加載： 使用torch.utils.data.distributed.DistributedSampler來包裝你的數據集，以確保每個進程處理不同的數據子集。

   from torch.utils.data import DataLoader, DistributedSampler
   
   dataset = YourDataset()
   sampler = DistributedSampler(dataset)
   dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=sampler)
   

6. 運行訓練： 現在你可以像平常一樣運行你的訓練腳本，PyTorch會自動處理分布式訓練的細節。

以下是一個完整的示例：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, DistributedSampler
import torchvision.datasets as datasets
import torchvision.transforms as transforms

# 初始化分布式環境
dist.init_process_group(backend='nccl', init_method='tcp://MASTER_IP:MASTER_PORT', world_size=WORLD_SIZE, rank=RANK)

# 定義模型
class YourModel(nn.Module):
    def __init__(self):
        super(YourModel, self).__init__()
        # 定義模型結構

    def forward(self, x):
        # 前向傳播
        return x

model = YourModel().to(rank)
ddp_model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank])

# 加載數據
transform = transforms.Compose([transforms.ToTensor()])
dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
sampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=sampler)

# 定義損失函數和優化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(ddp_model.parameters(), lr=0.01)

# 訓練模型
for epoch in range(EPOCHS):
    sampler.set_epoch(epoch)
    running_loss = 0.0
    for i, data in enumerate(dataloader, 0):
        inputs, labels = data[0].to(rank), data[1].to(rank)
        optimizer.zero_grad()
        outputs = ddp_model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch {epoch + 1}, Loss: {running_loss / len(dataloader)}')

通過以上步驟，你可以在Linux環境下使用PyTorch進行分布式訓練。