在CentOS上如何進行PyTorch的分布式訓練

centos系統上執行pytorch分布式訓練時,請按照以下流程操作:

  1. 安裝PyTorch:首先確認已安裝PyTorch。可以從PyTorch官方網站(https://www.php.cn/link/419e4410da152c74d727270283cb94ce

  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.distributed as dist  from torch.nn.parallel import DistributedDataParallel as DDP   def main():      # 初始化分布式架構      dist.init_process_group(backend='nccl', init_method='tcp://localhost:12345', world_size=4, rank=0)       # 定義模型并遷移至GPU      model = ... # 構建您的模型      model.cuda()       # 利用DistributedDataParallel封裝模型      model = DDP(model, device_ids=[torch.cuda.current_device()])       # 設置損失函數與優化器      criterion = torch.nn.CrossEntropyLoss().cuda()      optimizer = torch.optim.SGD(model.parameters(), lr=0.01)       # 載入數據      dataset = ... # 創建數據集      sampler = torch.utils.data.distributed.DistributedSampler(dataset)      dataloader = torch.utils.data.DataLoader(dataset, batch_size=..., sampler=sampler)       # 開始模型訓練      for epoch in range(...):          sampler.set_epoch(epoch)          for inputs, targets in dataloader:              inputs, targets = inputs.cuda(), targets.cuda()              optimizer.zero_grad()              outputs = model(inputs)              loss = criterion(outputs, targets)              loss.backward()              optimizer.step()       # 結束分布式架構      dist.destroy_process_group()   if __name__ == "__main__":      main()

    請依據實際情形調整模型、數據集、損失函數、優化器及訓練邏輯。

  4. 運行分布式訓練:可以借助mpirun或torch.distributed.launch來發起分布式訓練任務。例如:

     mpirun -np 4 python your_training_script.py

    或者采用torch.distributed.launch:

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

    其中的-np 4和–nproc_per_node=4指示每個節點所用GPU的數量。

  5. 關鍵點提示

    • 確保各節點間可通過網絡互相通信。
    • 核實所有節點的PyTorch版本與CUDA版本保持一致。
    • 若在多臺設備上實施分布式訓練,務必把MASTER_ADDR設為主機IP,并保證所有節點能通過此IP相互通信。

上述過程構成了一個基礎模板,具體應用時可能還需進一步定制化調整。開始分布式訓練前,推薦深入研讀PyTorch官方文檔中有關分布式訓練的部分。

? 版權聲明
THE END
喜歡就支持一下吧
點贊8 分享