Ubuntu PyTorch如何進行分布式訓練

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

1. 安裝PyTorch： 確保你已經安裝了PyTorch。你可以從PyTorch官網（https://pytorch.org/）獲取適合你系統的安裝命令。

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

   export CUDA_VISIBLE_DEVICES=0,1,2,3
   

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
   from torchvision import datasets, transforms
   
   # 初始化進程組
   torch.distributed.init_process_group(backend='nccl')
   
   # 創建模型并移動到GPU
   model = nn.Sequential(
       nn.Linear(784, 1024),
       nn.ReLU(),
       nn.Linear(1024, 10)
   ).to(torch.device("cuda"))
   
   # 使用DistributedDataParallel包裝模型
   model = DDP(model)
   
   # 定義損失函數和優化器
   criterion = nn.CrossEntropyLoss()
   optimizer = optim.SGD(model.parameters(), lr=0.01)
   
   # 加載數據集
   transform = transforms.Compose([transforms.ToTensor()])
   train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
   train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
   train_loader = DataLoader(dataset=train_dataset, batch_size=64, sampler=train_sampler)
   
   # 訓練模型
   for epoch in range(10):
       train_sampler.set_epoch(epoch)
       running_loss = 0.0
       for i, data in enumerate(train_loader, 0):
           inputs, labels = data[0].to(torch.device("cuda")), data[1].to(torch.device("cuda"))
           optimizer.zero_grad()
           outputs = model(inputs.view(-1, 784))
           loss = criterion(outputs, labels)
           loss.backward()
           optimizer.step()
           running_loss += loss.item()
       print(f'Epoch {epoch + 1}, Loss: {running_loss / len(train_loader)}')
   
   # 清理進程組
   torch.distributed.destroy_process_group()
   

4. 運行分布式訓練腳本： 使用torch.distributed.launch或accelerate庫來啟動分布式訓練。例如，如果你有4個GPU，你可以這樣運行腳本：

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

   或者，如果你安裝了accelerate庫，你可以這樣運行：

   accelerate launch your_training_script.py
   

5. 注意事項：

   • 確保所有節點上的PyTorch版本一致。
   • 確保所有節點都可以通過網絡互相通信。
   • 確保所有節點上的CUDA版本一致。
   • 如果你在云上運行，確保你已經配置了正確的安全組和網絡設置。

以上步驟提供了一個基本的分布式訓練框架，你可以根據自己的需求進行調整和優化。