pytorch同时使用多个gpu

在PyTorch中，可以通过使用torch.nn.DataParallel模块来同时在多个GPU上进行训练。这个模块可以自动将数据划分到不同的GPU上，并将模型复制到每个GPU上进行并行计算。要使用多个GPU，首先需要在代码中导入torch.nn.DataParallel模块。然后，将模型包装在这个模块中，如下所示：

import torch
import torch.nn as nn
from torch.nn import DataParallel

# 定义模型
model = MyModel()

# 将模型包装成DataParallel
model = DataParallel(model)

# 将模型移动到GPU
model = model.to(device)

接下来，在进行训练时，只需要将数据和模型放到GPU上即可，DataParallel会自动处理多个GPU的并行计算。例如：

# 将数据移动到GPU
input = input.to(device)
target = target.to(device)

# 前向计算
output = model(input)

# 后向传播和优化
loss.backward()
optimizer.step()

同时，要注意的是，确保torch.distributed.init_process_group()在使用多个GPU之前进行初始化，以确保多个进程间的通信和同步。

需要注意的是，并行训练需要确保多个GPU具有相同的显存大小，以及在使用多个GPU之前，需要确保每个GPU上的数据和模型都已经被正确地移动到对应的GPU上。