CentOS上PyTorch并行計算實踐

在CentOS上進行PyTorch并行計算實踐，通常涉及使用DistributedDataParallel（DDP）進行多GPU訓練。以下是一個簡單的實踐教程，展示了如何在CentOS系統上使用PyTorch進行分布式數據并行訓練。

環境準備

首先，確保你的CentOS系統上已經安裝了PyTorch和CUDA。你可以使用以下命令安裝PyTorch：

pip install torch torchvision torchaudio

示例代碼

以下是一個使用DistributedDataParallel進行并行訓練的簡單示例：

import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP

# 初始化進程組
def setup(rank, world_size):
    dist.init_process_group('nccl', rank=rank, world_size=world_size)

# 清理進程組
def cleanup():
    dist.destroy_process_group()

# 定義簡單的模型
class ToyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(1, 1)

    def forward(self, x):
        return self.layer(x)

# 自定義數據集
class MyDataset(Dataset):
    def __init__(self):
        self.data = torch.tensor([1, 2, 3, 4], dtype=torch.float32)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        return self.data[index:index + 1]

# 主函數
def main(rank, world_size):
    setup(rank, world_size)
    local_rank = rank
    device_id = local_rank % torch.cuda.device_count()

    dataset = MyDataset()
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=2, sampler=sampler)

    model = ToyModel().to(device_id)
    ddp_model = DDP(model, device_ids=[device_id])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    if rank == 0:
        torch.save(ddp_model.state_dict(), 'tmp.pth')

    for epoch in range(2):
        sampler.set_epoch(epoch)
        for data, target in dataloader:
            data, target = data.to(device_id), target.to(device_id)
            optimizer.zero_grad()
            output = ddp_model(data)
            loss = loss_fn(output, target)
            loss.backward()
            optimizer.step()

    cleanup()

if __name__ == '__main__':
    world_size = 4  # 假設有4張卡
    torch.multiprocessing.spawn(main, args=(world_size,), nprocs=world_size, join=True)

運行命令

假設有4張卡，使用以下命令在第三和第四張卡上并行運行：

export CUDA_VISIBLE_DEVICES=2,3
torchrun --nproc_per_node=2 dldemos/pytorchdistributed/main.py

注意事項

1. 環境變量：確保CUDA_VISIBLE_DEVICES環境變量正確設置，以便只使用指定的GPU。
2. 分布式訓練：使用DistributedDataParallel時，需要在多個進程間進行通信和同步，因此需要確保所有進程都在同一網絡中。
3. 資源管理：在訓練結束后，記得清理進程組以釋放資源。

通過以上步驟，你可以在CentOS上使用PyTorch進行分布式數據并行訓練。這種方法可以顯著提高訓練速度，特別是在處理大規模數據集和復雜模型時。