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进行分布式数据并行训练。这种方法可以显著提高训练速度，特别是在处理大规模数据集和复杂模型时。