Linux PyTorch如何進行分布式訓練

在Linux環境下使用PyTorch進行分布式訓練，可以遵循以下步驟：

環境準備

1. 安裝PyTorch： 確保你已經安裝了PyTorch，并且版本支持分布式訓練。

2. 設置環境變量： 為了啟用分布式訓練，需要設置一些環境變量，例如NCCL_DEBUG=INFO和HOROVOD_TIMELINE。

3. 網絡配置： 確保所有參與分布式訓練的節點可以通過SSH無密碼登錄，并且網絡配置正確。

啟動分布式訓練

1. 使用torch.distributed.launch： PyTorch提供了一個腳本torch.distributed.launch來啟動分布式訓練。這個腳本會自動處理一些初始化工作。

   python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE YOUR_TRAINING_SCRIPT.py
   

   其中，NUM_GPUS_YOU_HAVE是你每臺機器上的GPU數量，YOUR_TRAINING_SCRIPT.py是你的訓練腳本。

2. 使用horovodrun： 如果你使用Horovod進行分布式訓練，可以使用horovodrun命令來啟動訓練。

   horovodrun -np NUM_CPUS_YOU_HAVE YOUR_TRAINING_SCRIPT.py
   

   其中，NUM_CPUS_YOU_HAVE是你每臺機器上的CPU數量。

訓練腳本修改

在你的訓練腳本中，需要進行以下修改：

1. 初始化分布式環境： 使用torch.distributed.init_process_group來初始化分布式環境。

   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>
   )
   

   其中，<master_ip>是主節點的IP地址，<master_port>是主節點的端口號，<world_size>是總的進程數，<rank>是當前進程的排名。

2. 數據并行： 使用torch.nn.parallel.DistributedDataParallel來包裝你的模型。

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

3. 數據加載器： 確保你的數據加載器支持分布式訓練?？梢允褂?code>torch.utils.data.distributed.DistributedSampler。

   from torch.utils.data import DataLoader, DistributedSampler
   
   train_dataset = YourDataset()
   train_sampler = DistributedSampler(train_dataset)
   train_loader = DataLoader(train_dataset, batch_size=your_batch_size, sampler=train_sampler)
   

示例代碼

以下是一個簡單的分布式訓練示例：

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])

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

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

# 訓練循環
for epoch in range(num_epochs):
    train_sampler.set_epoch(epoch)
    running_loss = 0.0
    for i, data in enumerate(train_loader, 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(train_loader)}')

# 清理分布式環境
dist.destroy_process_group()

通過以上步驟，你可以在Linux環境下使用PyTorch進行分布式訓練。確保所有節點的網絡配置正確，并且環境變量設置無誤。