Ubuntu上如何進行PyTorch分布式訓練

在Ubuntu上進行PyTorch分布式訓練，你需要遵循以下步驟：

1. 安裝PyTorch： 首先，確保你已經安裝了PyTorch。你可以從PyTorch官網根據你的系統配置選擇合適的安裝命令。

2. 設置環境變量： 為了使用多GPU進行訓練，你需要設置一些環境變量。例如，如果你有4個GPU，你可以這樣設置：

   export CUDA_VISIBLE_DEVICES=0,1,2,3
   

   這將限制PyTorch只能看到和使用這四個GPU。

3. 編寫分布式訓練代碼： PyTorch提供了torch.nn.parallel.DistributedDataParallel類來幫助進行分布式訓練。你需要對你的模型和訓練循環進行一些修改來使用它。

   下面是一個簡單的例子，展示了如何使用DistributedDataParallel：

   import torch
   import torch.nn as nn
   import torch.optim as optim
   from torch.nn.parallel import DistributedDataParallel as DDP
   from torch.utils.data import DataLoader
   import torchvision.datasets as datasets
   import torchvision.transforms as transforms
   
   # 初始化進程組
   world_size = 4  # 假設你有4個GPU
   rank = ...  # 當前進程的rank（從0開始）
   torch.distributed.init_process_group(
       backend='nccl',  # 'nccl'是用于多GPU訓練的推薦后端
       init_method='env://',  # 使用環境變量來初始化
       world_size=world_size,
       rank=rank
   )
   
   # 創建模型并移動到對應的GPU
   model = ...  # 你的模型
   model.cuda(rank)
   
   # 使用DistributedDataParallel包裝模型
   model = DDP(model, device_ids=[rank])
   
   # 定義損失函數和優化器
   criterion = nn.CrossEntropyLoss().cuda(rank)
   optimizer = optim.SGD(model.parameters(), lr=0.01)
   
   # 加載數據集
   transform = transforms.Compose([transforms.ToTensor()])
   dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
   sampler = torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=world_size, rank=rank)
   dataloader = DataLoader(dataset, batch_size=64, sampler=sampler)
   
   # 訓練循環
   for epoch in range(num_epochs):
       sampler.set_epoch(epoch)  # 設置sampler的epoch，確保每個進程加載不同的數據
       for data, target in dataloader:
           data, target = data.cuda(rank), target.cuda(rank)
           optimizer.zero_grad()
           output = model(data)
           loss = criterion(output, target)
           loss.backward()
           optimizer.step()
       print(f'Epoch {epoch+1}, Loss: {loss.item()}')
   
   # 清理
   torch.distributed.destroy_process_group()
   

4. 運行分布式訓練： 使用torch.distributed.launch或accelerate庫來啟動分布式訓練。例如，如果你使用torch.distributed.launch，你的命令可能看起來像這樣：

   python -m torch.distributed.launch --nproc_per_node=4 your_training_script.py
   

   這里的--nproc_per_node=4指定了每個節點上使用的GPU數量。

5. 注意事項：

   • 確保所有節點都可以通過網絡相互通信。
   • 所有節點上的PyTorch版本和CUDA版本必須一致。
   • 分布式訓練通常需要更多的調試和測試來確保一切正常工作。

以上步驟提供了一個基本的框架，你可能需要根據你的具體情況進行調整。分布式訓練可能會涉及到更復雜的網絡配置和性能優化。