pytorch 使用多个gpu训练

PyTorch 提供了分布式训练功能，使得模型能够在多个 GPU 或者多台机器上并行运行，极大地提高了训练速度。为了利用多个 GPU 来加速训练，你可以使用以下步骤：

1. **安装必要的库**：首先，你需要安装 torch.distributed 和 torch.nn.parallel 库，这两个包分别用于分布式训练和数据并行。

   pip install torch torchvision torchaudio -f https://download.pytorch.org/whl/torch_stable.html

2. **初始化进程组**：创建一个 `DistributedDataParallel` 对象，这将帮你管理和同步多个 GPU 的梯度。需要设置好世界大小（number of processes），当前设备ID等信息。

   import torch.distributed as dist
   import torch.nn as nn

   if not is_main_process():
       # 这些操作仅主进程执行
       return

   dist.init_process_group(backend="nccl")  # 使用NCCL通信库，适用于GPU环境
   device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
   model = MyModel().to(device)
   model = nn.parallel.DistributedDataParallel(model)

3. **数据并行**：将数据集拆分成每个 GPU 能够处理的部分，并使用 DataLoader 并行加载。同时，在每个迭代中，确保模型接收的是分布到各个 GPU 的部分数据。

4. **更新权重**：在每个 mini-batch 训练完成后，使用 `model.module.parameters()` 获取模型的参数而不是直接的 `.parameters()`，因为 DistributedDataParallel 包装了参数。

5. **结束训练**：在训练完成时，通过 `dist.destroy_process_group()` 清理进程组。

for epoch in range(num_epochs):
    for data, target in train_loader:
        data, target = data.to(device), target.to(device)
        output = model(data)
        loss.backward()
        # 确保所有GPU更新完后再做优化
        optimizer.step()
        optimizer.zero_grad()

# 结束训练
if is_main_process():
    dist.destroy_process_group()