Linux PyTorch如何进行分布式训练

在Linux环境下使用PyTorch进行分布式训练，可以遵循以下步骤：

环境准备

1. 安装PyTorch： 确保你已经安装了PyTorch，并且版本支持分布式训练。

2. 设置环境变量： 为了启用分布式训练，需要设置一些环境变量，例如NCCL_DEBUG=INFO和HOROVOD_TIMELINE。

3. 网络配置： 确保所有参与分布式训练的节点可以通过SSH无密码登录，并且网络配置正确。

启动分布式训练

1. 使用torch.distributed.launch： PyTorch提供了一个脚本torch.distributed.launch来启动分布式训练。这个脚本会自动处理一些初始化工作。

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

   其中，NUM_GPUS_YOU_HAVE是你每台机器上的GPU数量，YOUR_TRAINING_SCRIPT.py是你的训练脚本。

2. 使用horovodrun： 如果你使用Horovod进行分布式训练，可以使用horovodrun命令来启动训练。

   horovodrun -np NUM_CPUS_YOU_HAVE YOUR_TRAINING_SCRIPT.py 

   其中，NUM_CPUS_YOU_HAVE是你每台机器上的CPU数量。

训练脚本修改

在你的训练脚本中，需要进行以下修改：

1. 初始化分布式环境： 使用torch.distributed.init_process_group来初始化分布式环境。

   import torch.distributed as dist dist.init_process_group( backend='nccl', # 或者 'gloo' init_method='tcp://<master_ip>:<master_port>', world_size=<world_size>, rank=<rank> ) 

   其中，<master_ip>是主节点的IP地址，<master_port>是主节点的端口号，<world_size>是总的进程数，<rank>是当前进程的排名。

2. 数据并行： 使用torch.nn.parallel.DistributedDataParallel来包装你的模型。

   model = YourModel().to(rank) ddp_model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank]) 

3. 数据加载器： 确保你的数据加载器支持分布式训练。可以使用torch.utils.data.distributed.DistributedSampler。

   from torch.utils.data import DataLoader, DistributedSampler train_dataset = YourDataset() train_sampler = DistributedSampler(train_dataset) train_loader = DataLoader(train_dataset, batch_size=your_batch_size, sampler=train_sampler) 

示例代码

以下是一个简单的分布式训练示例：

import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, DistributedSampler import torchvision.datasets as datasets import torchvision.transforms as transforms # 初始化分布式环境 dist.init_process_group( backend='nccl', init_method='tcp://<master_ip>:<master_port>', world_size=<world_size>, rank=<rank> ) # 定义模型 class YourModel(nn.Module): def __init__(self): super(YourModel, self).__init__() # 定义你的模型结构 def forward(self, x): # 定义前向传播 return x model = YourModel().to(rank) ddp_model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank]) # 定义损失函数和优化器 criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(ddp_model.parameters(), lr=0.01) # 加载数据 transform = transforms.Compose([transforms.ToTensor()]) train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform) train_sampler = DistributedSampler(train_dataset) train_loader = DataLoader(train_dataset, batch_size=64, sampler=train_sampler) # 训练循环 for epoch in range(num_epochs): train_sampler.set_epoch(epoch) running_loss = 0.0 for i, data in enumerate(train_loader, 0): inputs, labels = data[0].to(rank), data[1].to(rank) optimizer.zero_grad() outputs = ddp_model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() print(f'Epoch {epoch + 1}, Loss: {running_loss / len(train_loader)}') # 清理分布式环境 dist.destroy_process_group() 

通过以上步骤，你可以在Linux环境下使用PyTorch进行分布式训练。确保所有节点的网络配置正确，并且环境变量设置无误。