nn.DataParallel 是 PyTorch 中的一个模块,用于在多个 GPU 上并行运行模型。当有多个 GPU 并且想要利用它们来加速训练或推理时,这个模块会非常有用。nn.DataParallel 通过对模型中的每个子模块进行复制,并将输入数据分割成多个部分,然后在每个 GPU 上并行处理这些部分来实现并行化。
使用 nn.DataParallel 的基本步骤是:
- 初始化模型:首先,你需要定义你的神经网络模型。
- 包装模型:然后,使用 nn.DataParallel来包装你的模型。你需要传递你的模型和一个设备列表作为参数。设备列表通常是一个包含你想要使用的 GPU 的标识符的列表(例如 [0, 1]表示使用第一个和第二个 GPU)。
- 将数据移动到 GPU:在将数据传递给模型之前,你需要确保数据已经被移动到了 GPU上。这可以通过使用 .to(device) 方法来实现,其中 device 是一个表示目标 GPU 的对象。
- 前向传播:现在你可以像平常一样调用模型进行前向传播,nn.DataParallel 会在后台自动处理并行化。
import torch
import torch.nn as nn
import torch.nn.parallel
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel,self).__init__()
self.fc = nn.Linear(10,10)
def forward(self,x):
return self.fc(x)
model = SimpleModel()
if torch.cuda.device_count()>1:
print(torch.cuda.device_count())
model = nn.DataParallel(model,device_ids=[0,1])
#将模型移动到GPU上
model.to("cuda")
input_data = torch.randn(16,10).to("cuda")
output_data = model(input_data)