pytorch 多GPU

PyTorch 是一种广泛使用的深度学习框架，它支持多 GPU 训练，这对于处理大型模型和加速训练过程非常有用。在 PyTorch 中，你可以利用 GPU 的并行计算能力来提高模型训练的速度。

使用多 GPU 的主要步骤包括：

1. **设置设备**：首先，你需要将模型放置到正确的设备上，如 `cuda()` 或 `cuda(device_id)`，其中 `device_id` 表示你想使用的 GPU 的编号。可以使用 `torch.cuda.device_count()` 来查看可用的 GPU 数量。

import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

2. **数据加载**：使用 DataLoader 可以轻松地从多个 GPU 加载数据。你可以设置 `batch_size` 进行批处理，并指定 `pin_memory=True` 避免内存复制。

3. **模型复制**：对于分布式训练，你需要创建多个模型副本并将它们放在不同的 GPU 上。通常通过 `nn.DataParallel` 或 `torch.nn.parallel.DistributedDataParallel` 类实现。

model = nn.DataParallel(model)
# 或者
if torch.distributed.is_initialized():
    model = DistributedDataParallel(model)

4. **同步梯度**：在多个 GPU 之间同步梯度是很重要的，这通常发生在优化循环结束时。PyTorch 提供了多种同步策略，比如 `torch.nn.parallel.AvgGradReducer` 和 `torch.optim.SynchronizedBatchNorm1d`。

5. **训练循环**：在训练过程中，每次迭代都会将数据分布在所有 GPU 上，然后更新模型参数。

for inputs, targets in dataloader:
    inputs, targets = inputs.to(device), targets.to(device)
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()