在深度学习的实践中,由于数据集非常大,模型训练需要很长时间,为了解决这个问题,我们可以利用多个GPU并行计算来加速训练过程。torch.nn.DataParallel() 是PyTorch提供的一个非常方便的工具,它可以让我们在多个GPU上并行训练模型。

1、DataParallel的原理

首先在指定的每张显卡上拷贝一份模型,然后将输入的数据分散到各张显卡上,计算梯度,然后回传到第一张显卡上,最后再对模型进行参数优化。

需要注意的是:第一张显卡的负载往往更高,但由于该方法集成度高,书写简便,使用仍十分广泛。

2、DataParallel参数介绍

torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

(1)module:表示我们自定义的模型;
(2)device_ids:表示我们用于训练的GPU device ids;
(3)output_device:这个参数表示输出结果的device。而这最后一个参数output_device一般情况下是省略不写的,那么默认就是在device_ids[0],也就是第一块卡上,也就解释了为什么第一块卡的显存会占用的比其他卡要更多一些。

更多介绍请参考官网:https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html

3、DataParallel应用举例

# 第一步:将模型移动到GPU设备上
model = model.cuda(0) 

# 第二步:设置显卡id
device_ids = [0, 1]
 
# 第三步:拷贝模型副本到GPU 0和 GPU 1,因为GPU 0已经存在,不需要拷贝
model = torch.nn.DataParallel(model, device_ids=device_ids)

# 第四步:准备数据。因为所有数据都需要先放到指定的第一张显卡上才能进行多卡训练
data = data.cuda(0)

# 第五步:训练模型
output = model(data)

4、DataParallel底层流程

第一步:数据分割。DataParallel会自动将数据分割成多个部分,每个部分都会在一个GPU上进行处理。分割的方式取决于输入数据的形状和GPU的数量。
第二步:创建模型副本:在每个GPU上,都会创建一个模型的副本。这些副本共享相同的参数,但每个副本都独立地处理一部分输入数据。
第三步:结果合并:在所有GPU上的处理完成后,DataParallel会将结果合并成一个完整的输出。这个过程是自动的,我们不需要手动进行合并。