pytorch多卡分布式训练

PyTorch支持多卡分布式训练，可以利用多个GPU进行模型的并行计算。下面是一个简单的示例代码，展示了如何使用PyTorch进行多卡分布式训练：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel

# 初始化进程组
dist.init_process_group(backend='nccl')

# 设置设备和参数
device = torch.device('cuda')
batch_size = 64
learning_rate = 0.01
num_epochs = 10

# 定义模型
model = nn.Linear(10, 1).to(device)
model = DistributedDataParallel(model)

# 定义数据加载器
train_dataset = ...
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler)

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

# 训练模型
for epoch in range(num_epochs):
    for inputs, labels in train_loader:
        inputs = inputs.to(device)
        labels = labels.to(device)

        outputs = model(inputs)
        loss = criterion(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# 清理进程组
dist.destroy_process_group()

在这个示例中，我们首先使用`dist.init_process_group`初始化进程组，然后将模型包装在`DistributedDataParallel`中以实现多卡并行计算。接下来，我们创建了一个数据加载器，并使用`DistributedSampler`进行数据的分布式采样。然后，我们定义了损失函数和优化器，并在每个epoch内进行训练。最后，我们使用`dist.destroy_process_group`清理进程组。

请注意，这只是一个简单的示例，实际的多卡分布式训练可能还涉及到其他内容，如数据并行、模型的同步等。具体的实现方式可能因具体的情况而有所不同，请根据自己的需求进行相应的调整。