使用pytorch写多GPU并行代码

在 PyTorch 中使用多 GPU 并行需要使用 nn.DataParallel 或 nn.parallel.DistributedDataParallel 类。

首先，您需要在初始化模型时将模型放在 GPU 上：

model = MyModel()
if torch.cuda.device_count() > 1:
  print("Using", torch.cuda.device_count(), "GPUs")
  model = nn.DataParallel(model)
model.to(device)

然后，您可以在训练循环中将输入数据和标签放在 GPU 上，并在 forward 函数中使用模型：

for input, label in training_data:
  input = input.to(device)
  label = label.to(device)
  output = model(input)
  loss = criterion(output, label)
  optimizer.zero_grad()
  loss.backward()
  optimizer.step()

如果您想使用分布式训练，则可以使用 nn.parallel.DistributedDataParallel 类。这需要使用 torch.nn.parallel.init_process_group 函数来初始化分布式环境。

torch.nn.parallel.init_process_group(backend='nccl')
model = MyModel()
model = nn.parallel.DistributedDataParallel(model)

然后，您可以使用与 nn.DataParallel 相同的方法在训练循环中使用模型。

请注意，您还需要使用 torch.nn.utils.data.DistributedSampler 类来对数据进行采样，以确保在分布式环境中正确地对数据进行划分。

train_sampler = torch.nn.utils.data.DistributedSampler(training_data)
train_loader = torch.utils.data.DataLoader(training_data, sampler=train_sampler, ...)

希望这能帮助您。