PyTorch中多GPU训练的利器——nn.DataParallel()

在深度学习中，使用多个GPU可以加快模型的训练速度和增加模型的容量。然而，手动编写多GPU训练代码是一项复杂而耗时的任务。为了简化这个过程，PyTorch提供了一个方便的工具——nn.DataParallel()，它可以自动将模型的输入和输出分配到多个GPU上，并在每个GPU上并行计算梯度。本文将详细介绍nn.DataParallel()的用法，并提供一个简单的使用例子。

首先，我们需要准备一个可以在多个GPU上并行计算的模型。以下是一个简单的示例：

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(100, 10)
    
    def forward(self, x):
        return self.fc(x)

在这个示例中，我们定义了一个包含一个全连接层的简单模型。接下来，我们需要将模型包裹在nn.DataParallel()中，以便在多个GPU上进行并行计算。以下是具体的代码：

model = SimpleModel()
model = nn.DataParallel(model)

通过这两行代码，我们将SimpleModel包裹在nn.DataParallel()中，并将结果赋给model。这样，model就成为了一个可以在多个GPU上并行计算的模型。

接下来，我们需要将模型和数据移到GPU上。假设我们有两个GPU，可以使用以下代码将模型和数据移到GPU上：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
inputs = inputs.to(device)

在这个示例中，我们首先检查是否有可用的GPU，如果有，则将device设置为 个GPU，否则设置为CPU。然后，我们使用model.to(device)将模型移到GPU上，使用inputs.to(device)将数据移到GPU上。

接下来，我们可以像在单个GPU上一样使用模型进行训练或推理。nn.DataParallel()会自动将输入数据分配到多个GPU上并行计算，并将输出结果合并成一个结果。下面是一个简单的使用例子：

for epoch in range(num_epochs):
    for inputs, targets in dataloader:
        inputs = inputs.to(device)
        targets = targets.to(device)
        
        outputs = model(inputs)
        
        # 这里是你的训练或推理代码

在这个示例中，我们使用一个简单的循环来遍历训练数据。首先，我们将inputs和targets移动到GPU上。然后，我们通过model(inputs)调用模型，nn.DataParallel()会自动将inputs分配到多个GPU上并行计算，并将结果合并成一个outputs。最后，我们可以使用outputs进行训练或推理。

使用nn.DataParallel()进行多GPU训练非常简单，只需要将模型包裹在nn.DataParallel()中，并将模型和数据移到GPU上即可。nn.DataParallel()会自动处理并行计算和梯度更新的细节，让你更专注于模型本身的开发和调试。