PyTorch在Ubuntu上的分布式训练怎么做 - 问答

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

安装PyTorch：首先，确保你已经安装了PyTorch。你可以从PyTorch官网（https://pytorch.org/）获取适合你系统的安装指令。
设置环境变量：为了使用多GPU进行分布式训练，你需要设置一些环境变量。例如，如果你有4个GPU，你可以这样设置：
```
export MASTER_ADDR='localhost'
export MASTER_PORT='12345'
export WORLD_SIZE=4
```
MASTER_ADDR 是主节点的IP地址，MASTER_PORT 是一个随机选择的端口号，WORLD_SIZE 是参与训练的GPU总数。

编写分布式训练脚本：在你的训练脚本中，你需要使用 torch.distributed.launch 或者 accelerate 库来启动分布式训练。以下是使用 torch.distributed.launch 的一个简单示例：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP

# 初始化进程组
torch.distributed.init_process_group(
    backend='nccl',  # 'nccl' is recommended for distributed GPU training
    init_method='tcp://<master_ip>:<master_port>',
    world_size=<world_size>,
    rank=<rank>
)

# 创建模型并将其移动到对应的GPU
model = nn.Linear(10, 10).to(rank)

# 使用DistributedDataParallel包装模型
ddp_model = DDP(model, device_ids=[rank])

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

# 假设data_loader是一个已经定义好的数据加载器
for data, target in data_loader:
    data, target = data.to(rank), target.to(rank)
    
    # 前向传播
    output = ddp_model(data)
    
    # 计算损失
    loss = criterion(output, target)
    
    # 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# 清理进程组
torch.distributed.destroy_process_group()

在这个脚本中，<master_ip> 和 <master_port> 需要替换为实际的主节点IP地址和端口号，<world_size> 是总的GPU数量，<rank> 是当前进程的排名（从0开始）。

运行分布式训练：使用 torch.distributed.launch 来运行你的训练脚本。例如，如果你的脚本名为 train.py，你可以这样运行它：
```
python -m torch.distributed.launch --nproc_per_node=<num_gpus_per_node> train.py
```
其中 <num_gpus_per_node> 是每个节点上的GPU数量。
使用 accelerate 库（可选）： accelerate 是一个由Hugging Face提供的库，它可以简化分布式训练的设置。使用 accelerate，你可以用更少的代码行来启动分布式训练。首先，你需要安装 accelerate：
```
pip install accelerate
```
然后，你可以使用 accelerate 的 launch 命令来运行你的训练脚本：
```
accelerate launch train.py
```
accelerate 会自动检测可用的GPU，并为你设置好分布式训练的环境。

请注意，这些步骤假设你已经有了一个可以运行的单机多GPU训练脚本。分布式训练通常涉及到更多的复杂性，比如数据并行化、模型并行化、梯度聚合等。确保你的代码在这些方面也是正确设置的。

0 赞

0 踩