PyTorch在Ubuntu上的分布式训练如何实现 - 问答

在Ubuntu上使用PyTorch进行分布式训练，通常涉及以下几个步骤：

安装PyTorch：确保你已经安装了PyTorch。你可以从PyTorch官网获取适合你系统的安装命令。
设置环境变量：为了使用多GPU训练，你需要设置一些环境变量，例如NCCL_DEBUG=INFO和NCCL_SOCKET_IFNAME。
编写分布式训练脚本：使用PyTorch的torch.nn.parallel.DistributedDataParallel来包装你的模型，并使用torch.distributed.launch或accelerate库来启动分布式训练。
运行分布式训练：在命令行中使用mpirun或python -m torch.distributed.launch来启动多个进程进行训练。

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

步骤 1: 安装PyTorch

pip install torch torchvision torchaudio

步骤 2: 设置环境变量

在运行训练脚本之前，设置必要的环境变量：

export NCCL_DEBUG=INFO
export NCCL_SOCKET_IFNAME=^docker0,lo  # 根据实际情况修改

步骤 3: 编写分布式训练脚本

创建一个名为distributed_train.py的脚本，内容如下：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
import torchvision.datasets as datasets
import torchvision.transforms as transforms

# 初始化进程组
world_size = 4  # 假设使用4个GPU
rank = 0  # 当前进程的排名
torch.distributed.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=rank)

# 定义模型
model = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10),
)
model.cuda(rank)  # 将模型放到对应的GPU上
model = DDP(model, device_ids=[rank])  # 使用DistributedDataParallel包装模型

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

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

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

# 清理
torch.distributed.destroy_process_group()

步骤 4: 运行分布式训练

使用torch.distributed.launch来启动分布式训练：

python -m torch.distributed.launch --nproc_per_node=4 distributed_train.py

这里的--nproc_per_node=4表示每个节点使用4个GPU。如果你有多个节点，你需要指定--nnodes、--node_rank和master_addr等参数。

注意：在实际部署时，你可能需要根据你的集群配置调整上述脚本和命令。此外，确保所有节点之间的网络连接是畅通的，并且防火墙设置允许所需的端口通信。

0 赞

0 踩