linux

Linux PyTorch如何进行分布式训练

小樊
46
2025-05-13 23:45:01
栏目: 智能运维

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

环境准备

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

  2. 设置环境变量: 为了启用分布式训练,需要设置一些环境变量,例如NCCL_DEBUG=INFOHOROVOD_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进行分布式训练。确保所有节点的网络配置正确,并且环境变量设置无误。

0
看了该问题的人还看了