在CentOS上進行PyTorch的分布式訓練，你需要遵循以下步驟：
安裝PyTorch： 確保你已經安裝了PyTorch。你可以從PyTorch官網根據你的系統和CUDA版本選擇合適的安裝命令。
準備環境： 在開始分布式訓練之前，確保所有參與訓練的節點都能夠通過網絡互相通信，并且可以SSH無密碼登錄。
編寫分布式訓練腳本：
PyTorch提供了torch.distributed包來支持分布式訓練。你需要編寫一個腳本來初始化分布式環境，設置模型、優化器等，并使用DistributedDataParallel來包裝你的模型。
啟動分布式訓練：
使用mpirun或torch.distributed.launch工具來啟動分布式訓練。你需要指定總的GPU數量、每個節點的GPU數量、節點地址和端口等信息。
下面是一個簡單的例子，展示了如何使用torch.distributed.launch來啟動分布式訓練：
# run.py
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import torchvision.models as models
def main(rank, world_size):
# 初始化進程組
dist.init_process_group(
backend='nccl', # 'nccl' is recommended for distributed GPU training
init_method='tcp://<master_ip>:<master_port>', # 替換為你的主節點IP和端口
world_size=world_size, # 總的進程數
rank=rank # 當前進程的排名
)
# 創建模型并移動到對應的GPU
model = models.resnet18(pretrained=True).to(rank)
# 使用DistributedDataParallel包裝模型
ddp_model = DDP(model, device_ids=[rank])
# 創建損失函數和優化器
criterion = torch.nn.CrossEntropyLoss().to(rank)
optimizer = torch.optim.SGD(ddp_model.parameters(), lr=0.01)
# 加載數據集并進行分布式采樣
dataset = ... # 你的數據集
sampler = torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=world_size, rank=rank)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=..., sampler=sampler)
# 訓練模型
for epoch in range(...): # 替換為你的epoch數
sampler.set_epoch(epoch)
for inputs, targets in dataloader:
inputs, targets = inputs.to(rank), targets.to(rank)
optimizer.zero_grad()
outputs = ddp_model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
# 清理進程組
dist.destroy_process_group()
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--world-size', type=int, default=4, help='number of distributed processes')
parser.add_argument('--rank', type=int, default=0, help='rank of the current process')
args = parser.parse_args()
main(args.rank, args.world_size)
啟動分布式訓練的命令可能如下所示：
mpirun -np 4 python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE run.py
或者使用torch.distributed.launch：
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --nnodes=NUM_NODES_YOU_HAVE --node_rank=NODE_RANK_YOU_HAVE --master_addr=MASTER_NODE_IP --master_port=MASTER_NODE_PORT run.py
在這里，NUM_GPUS_YOU_HAVE是你每個節點上的GPU數量，NUM_NODES_YOU_HAVE是節點總數，NODE_RANK_YOU_HAVE是當前節點的排名（從0開始），MASTER_NODE_IP是主節點的IP地址，MASTER_NODE_PORT是主節點上用于通信的端口號。
請注意，這只是一個基本的例子，實際的分布式訓練腳本可能需要更多的配置和優化。此外，確保你的網絡設置允許節點間的通信，并且防火墻規則不會阻止必要的端口。
