在这个教程里,我们将学习如何使用数据并行(DataParallel
)来使用多GPU。
PyTorch非常容易的就可以使用GPU,可以用如下方式把一个模型放到GPU上:
device = torch.device("cuda:0")
model.to(device)
然后可以复制所有的张量到GPU上:
mytensor = my_tensor.to(device)
请注意,调用my_tensor.to(device)
返回一个GPU上的my_tensor
副本,而不是重写my_tensor
。我们需要把它赋值给一个新的张量并在GPU上使用这个张量。
在多GPU上执行前向和反向传播是自然而然的事。然而,PyTorch默认将只是用一个GPU。你可以使用DataParallel
让模型并行运行来轻易的让你的操作在多个GPU上运行。
model = nn.DataParallel(model)
这是这篇教程背后的核心,我们接下来将更详细的介绍它。
导入PyTorch模块和定义参数。