pytorch ddp 多机多卡训练示例

下面是一个使用PyTorch DDP（分布式数据并行）进行多机多卡训练的示例：
1.首先，需要使用torch.distributed.launch启动多个进程，每个进程都运行相同的脚本并使用不同的参数。例如，在两台机器上运行以下命令：
# Machine 1
python -m torch.distributed.launch --nproc_per_node=2 --nnodes=2 --node_rank=0 --master_addr="10.0.0.1" --master_port=8888 train.py

# Machine 2
python -m torch.distributed.launch --nproc_per_node=2 --nnodes=2 --node_rank=1 --master_addr="10.0.0.1" --master_port=8888 train.py
登录后复制
上面的命令将在两台机器上启动4个进程，每个进程使用2个GPU进行训练。
2.在代码中，使用torch.distributed初始化进程组，并将模型和数据加载到每个进程中。例如：
import torch
import torch.nn as nn
import torch.distributed as dist

# Initialize distributed process group
dist.init_process_group(backend='nccl', init_method='env://')

# Load data and model
train_data = ...
train_loader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True)

model = nn.Sequential(
nn.Linear(784, 256),
nn.ReLU(),
nn.Linear(256, 10)
)

# Distributed model and optimizer
model = nn.parallel.DistributedDataParallel(model)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
登录后复制
这里使用了nn.parallel.DistributedDataParallel将模型包装成分布式模型，使用torch.optim.SGD作为优化器。
3.在训练循环中，每个进程都会收集自己的梯度并将它们聚合到进程组中。然后，所有进程都将使用平均梯度更新模型参数。例如：
for epoch in range(10):
for data, target in train_loader:
optimizer.zero_grad()
output = model(data)
loss = nn.functional.cross_entropy(output, target)
loss.backward()

# All-reduce gradients
for param in model.parameters():
dist.all_reduce(param.grad.data, op=dist.reduce_op.SUM)

optimizer.step()
登录后复制
在每个批次之后，使用dist.all_reduce将每个进程的梯度聚合到进程组中，然后使用平均梯度更新模型参数。
4.训练完成后，使用dist.destroy_process_group()关闭进程组并释放资源。例如：
dist.destroy_process_group()
登录后复制
这个示例展示了如何使用PyTorch DDP进行多机多卡训练。需要注意的是，使用DDP需要确保所有进程都能够访问相同的数据和模型，并且需要正确设置进程组中的参数。