【PyTorch教程】如何使用PyTorch分布式并行模块DistributedDataParallel(DDP)进行多卡训练

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
 
# 假设已经初始化了进程组，例如在主进程中使用
# torch.distributed.init_process_group(backend, init_method, world_size, rank)
 
# 定义一个简单的模型
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc = nn.Linear(10, 10)
    
    def forward(self, x):
        return self.fc(x)
 
model = Model()
 
# 在每个进程中，创建模型的DDP实例
# 需要设置device_ids参数为当前进程的rank
model = DDP(model, device_ids=[dist.get_rank()])
 
# 定义优化器和损失函数
optimizer = optim.SGD(model.parameters(), lr=0.001)
loss_fn = nn.MSELoss()
 
# 假设已经有数据和标签，进行正向传播、计算损失和反向传播
inputs = ...
labels = ...
 
optimizer.zero_grad()
outputs = model(inputs)
loss_fn(outputs, labels).backward()
optimizer.step()
 
# 注意：在实际使用时，还需要配合torch.distributed.launch或其他集群管理工具使用。

这段代码展示了如何在PyTorch中使用DistributedDataParallel进行模型训练。在实际应用中，你需要确保分布式环境已经正确设置，并且使用合适的启动脚本来启动你的进程，例如torch.distributed.launch。