Linux上PyTorch如何进行分布式训练 - 问答

在Linux上进行PyTorch的分布式训练，你需要遵循以下步骤：

环境准备：
- 确保你有一个支持分布式训练的环境，通常需要多个GPU或者多台机器。
- 安装PyTorch，确保它支持分布式训练。你可以从PyTorch官网获取适合你系统的安装命令。
初始化分布式环境：
- 使用torch.distributed.init_process_group函数来初始化分布式环境。你需要指定后端（如nccl、gloo等）、初始化方法（如tcp://或intra）、世界大小（即总的进程数）和当前进程的rank（每个进程的唯一标识符）。
数据并行：
- 使用torch.nn.parallel.DistributedDataParallel（DDP）来包装你的模型。DDP会自动处理数据的划分、梯度的聚合等。
编写分布式训练代码：
- 在你的训练脚本中，需要添加逻辑来处理分布式训练的特殊部分，比如数据加载器的设置、模型的分发、梯度的同步等。
运行分布式训练：
- 使用mpirun、mpiexec或者torch.distributed.launch来启动分布式训练。你需要指定总的进程数、每个进程的入口脚本等。

下面是一个简单的例子，展示了如何在Linux上使用PyTorch进行分布式训练：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist
from torch.utils.data import DataLoader, DistributedSampler
import torchvision.datasets as datasets
import torchvision.transforms as transforms

def main(rank, world_size):
    # 初始化分布式环境
    dist.init_process_group(
        backend='nccl',
        init_method='tcp://localhost:23456',
        world_size=world_size,
        rank=rank
    )

    # 创建模型并将其移动到对应的GPU
    model = ... # 定义你的模型
    model.cuda(rank)
    model = DDP(model, device_ids=[rank])

    # 创建损失函数和优化器
    criterion = nn.CrossEntropyLoss().cuda(rank)
    optimizer = optim.SGD(model.parameters(), lr=0.01)

    # 加载数据集
    transform = transforms.Compose([transforms.ToTensor()])
    dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=64, sampler=sampler)

    # 训练模型
    for epoch in range(10):
        sampler.set_epoch(epoch)
        running_loss = 0.0
        for i, data in enumerate(dataloader, 0):
            inputs, labels = data[0].cuda(rank), data[1].cuda(rank)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        print(f'Rank {rank}, Epoch {epoch}, Loss: {running_loss/len(dataloader)}')

    # 清理分布式环境
    dist.destroy_process_group()

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--world_size', type=int, default=2)
    parser.add_argument('--rank', type=int, default=0)
    args = parser.parse_args()

    main(args.rank, args.world_size)

要运行这个脚本，你需要使用torch.distributed.launch或者mpirun/mpiexec。例如：

python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE YOUR_TRAINING_SCRIPT.py

或者使用mpirun：

mpirun -np NUM_GPUS_YOU_HAVE -bind-to none -map-by slot -x NCCL_DEBUG=INFO -x LD_LIBRARY_PATH -x PATH -mca pml ob1 -mca btl ^openib YOUR_TRAINING_SCRIPT.py

请注意，这里的NUM_GPUS_YOU_HAVE是你想要使用的GPU数量，YOUR_TRAINING_SCRIPT.py是你的训练脚本的名称。在分布式训练中，每个进程都会运行相同的脚本，但是它们会有不同的rank和world_size参数。

0 赞

0 踩