PyTorch在CentOS中的分布式訓練實踐

在CentOS上進行PyTorch的分布式訓練，可以參考以下步驟：

環境準備

1. 安裝Anaconda或Miniconda：首先，在CentOS上安裝Anaconda或Miniconda。這里以安裝Miniconda為例：

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh

按照提示完成安裝過程。

2. 創建虛擬環境：創建一個用于PyTorch開發的虛擬環境，并激活它：

conda create -n pytorch_env python=3.10
conda activate pytorch_env

3. 配置鏡像源（可選）：為了加快下載速度，可以配置國內鏡像源。編輯.condarc文件（通常位于用戶主目錄下）：

channels:
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/msys2/
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forges/
show_channel_urls: true
auto_activate_base: false

4. 安裝PyTorch：根據你的CUDA版本選擇合適的PyTorch版本進行安裝。以下是一些常見的安裝命令：

• 安裝CPU版本：

pip install torch torchvision torchaudio cpuonly -c pytorch

• 安裝GPU版本（需要CUDA支持）：

nvidia-smi  # 查看CUDA版本
pip install torch torchvision torchaudio cudatoolkit=<your_cuda_version> -c pytorch

例如，對于CUDA 12.1：

pip install torch torchvision torchaudio cudatoolkit=12.1 -c pytorch

5. 驗證安裝：安裝完成后，可以驗證PyTorch是否安裝成功：

import torch
print(torch.__version__)
print(torch.cuda.is_available())

如果沒有報錯，并且torch.cuda.is_available()返回True，則表示安裝成功。

編寫分布式訓練腳本

使用PyTorch的torch.distributed包來編寫分布式訓練腳本。你需要使用torch.nn.parallel.DistributedDataParallel來包裝你的模型，并使用torch.distributed.launch或者mpirun來啟動分布式訓練。

以下是一個簡單的分布式訓練腳本示例：

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.multiprocessing as mp
import torchvision.models as models

def main(rank, world_size):
    dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=rank)
    model = models.resnet18(pretrained=True).to(rank)
    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=256, sampler=sampler)
    
    for epoch in range(10):  # 替換為你的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()
            print(f'Rank {rank}, Epoch {epoch}, Loss {loss.item()}')
    
    dist.destroy_process_group()

if __name__ == '__main__':
    world_size = 4  # 假設有4個GPU
    mp.spawn(main, args=(world_size,), nprocs=world_size, join=True)

啟動分布式訓練

使用mpirun或torch.distributed.launch來啟動分布式訓練。例如，如果你想在兩個GPU上運行訓練腳本，可以使用以下命令：

mpirun -np 2 python your_training_script.py

或者使用torch.distributed.launch：

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

監控和調試

分布式訓練可能會遇到各種問題，包括網絡通信問題、同步問題等。使用nccl-tests來測試你的GPU之間的通信是否正常。同時，確保你的日志記錄是詳細的，以便于調試。

以上步驟應該能幫助你在CentOS上成功進行PyTorch的分布式訓練。如果在安裝過程中遇到問題，可以參考PyTorch官方文檔或相關社區論壇尋求幫助。