Ubuntu上PyTorch的分布式訓練

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

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

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

   export MASTER_ADDR='localhost'
   export MASTER_PORT='12345'
   export WORLD_SIZE=4
   

   MASTER_ADDR 是主節點的IP地址，MASTER_PORT 是一個未被使用的端口號，WORLD_SIZE 是參與訓練的GPU總數。

3. 編寫分布式訓練腳本： 在你的PyTorch腳本中，你需要使用torch.distributed包來初始化分布式環境。以下是一個簡單的例子：

   import torch
   import torch.nn as nn
   import torch.optim as optim
   from torch.nn.parallel import DistributedDataParallel as DDP
   
   # 初始化進程組
   torch.distributed.init_process_group(
       backend='nccl',  # 'nccl' is recommended for distributed GPU training
       init_method='tcp://localhost:12345',
       world_size=4,
       rank=0  # 這個rank需要在每個進程中設置為不同的值
   )
   
   # 創建模型并將其移動到GPU
   model = nn.Linear(10, 10).to(torch.device("cuda"))
   
   # 使用DistributedDataParallel包裝模型
   ddp_model = DDP(model, device_ids=[rank])
   
   # 創建損失函數和優化器
   criterion = nn.CrossEntropyLoss()
   optimizer = optim.SGD(ddp_model.parameters(), lr=0.01)
   
   # 假設我們有一些數據
   inputs = torch.randn(20, 10).to(torch.device("cuda"))
   labels = torch.randint(0, 10, (20,)).to(torch.device("cuda"))
   
   # 訓練循環
   for epoch in range(10):
       optimizer.zero_grad()
       outputs = ddp_model(inputs)
       loss = criterion(outputs, labels)
       loss.backward()
       optimizer.step()
       print(f"Rank {rank}, Epoch {epoch}, Loss {loss.item()}")
   

4. 運行分布式訓練： 使用torch.distributed.launch工具來啟動分布式訓練。例如，如果你有一個名為train.py的訓練腳本，你可以這樣運行它：

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

   --nproc_per_node 參數指定了每個節點上使用的GPU數量。

請注意，這些步驟假設你已經有了一個可以工作的PyTorch環境，并且你的系統中有至少4個GPU可用。分布式訓練可能會涉及到更復雜的網絡配置和同步問題，因此在實際應用中可能需要更多的設置和調試。