pytorch模型训练单机多卡(二):数据并行DataParallel

1.并行处理机制

DataParallel系统通过将整个小型批处理加载到主线程上,然后将子小型批处理分散到整个GPU网络中来工作。

具体是将输入一个 batch 的数据均分成多份,分别送到对应的 GPU 进行计算。与 Module 相关的所有数据也都会以浅复制的方式复制多份。每个 GPU 在单独的线程上将针对各自的输入数据独立并行地进行 forward 计算。然后在主GPU上收集网络输出,并通过将网络输出与批次中每个元素的真实数据标签进行比较来计算损失函数值。接下来,损失值分散给各个GPU,每个GPU进行反向传播以计算梯度。最后,在主GPU上归约梯度、进行梯度下降,并更新主GPU上的模型参数。由于模型参数仅在主GPU上更新,而其他从属GPU此时并不是同步更新的,所以需要将更新后的模型参数复制到剩余的从属 GPU 中,以此来实现并行。

DataParallel会将定义的网络模型参数默认放在GPU 0上,所以dataparallel实质是可以看做把训练参数从GPU拷贝到其他的GPU同时训练,这样会导致内存和GPU使用率出现很严重的负载不均衡现象,即GPU 0的使用内存和使用率会大大超出其他显卡的使用内存,因为在这里GPU0作为master来进行梯度的汇总和模型的更新,再将计算任务下发给其他GPU,所以他的内存和使用率会比其他的高。

2.流程图

3.pytorch代码

非常简单:

device_ids = [0, 1]
net = torch.nn.DataParallel(net, device_ids=device_ids)

但是,模型的保存、更新和调用也需要配套。否则bug层出不穷。

使用nn.DataParallel后,事实上DataParallel也是一个Pytorch的nn.Module,那么你的模型和优化器都需要使用.module来得到实际的模型和优化器,如下:

保存模型:
torch.save(net.module.state_dict(), path)
加载模型:
net=nn.DataParallel(Resnet18())
net.load_state_dict(torch.load(path))
net=net.module
优化器使用:
optimizer.step() --> optimizer.module.step()

还有一个问题就是,如果直接使用nn.DataParallel的时候,训练采用多卡训练,会出现一个warning:

UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; 
will instead unsqueeze and return a vector.

 首先说明一下:每张卡上的loss都是要汇总到第0张卡上求梯度,更新好以后把权重分发到其余卡。但是为什么会出现这个warning,这其实和nn.DataParallel中最后一个参数dim有关,其表示tensors被分散的维度,默认是0,nn.DataParallel将在dim0(批处理维度)中对数据进行分块,并将每个分块发送到相应的设备。单卡的没有这个warning,多卡的时候采用nn.DataParallel训练会出现这个warning,由于计算loss的时候是分别在多卡计算的,那么返回的也就是多个loss,你使用了多少个gpu,就会返回多少个loss。(有人建议DataParallel类应该有reduce和size_average参数,比如用于聚合输出的不同loss函数,最终返回一个向量,有多少个gpu,返回的向量就有几维。)

关于这个问题在pytorch官网的issues上有过讨论:https://github.com/pytorch/pytorch/issues/9811

前期探讨中,有人提出求loss平均的方式会在不同数量的gpu上训练会以微妙的方式影响结果。模块返回该batch中所有损失的平均值,如果在4个gpu上运行,将返回4个平均值的向量。然后取这个向量的平均值。但是,如果在3个GPU或单个GPU上运行,这将不是同一个数字,因为每个GPU处理的batch size不同!

似乎一看,这么求平均loss确实有不合理的地方。那么有什么好的解决办法呢,可以使用size_average=False,reduce=True作为参数。每个GPU上的损失将相加,但不除以GPU上的批大小。然后将所有平行损耗相加,除以整批的大小,那么不管几块GPU最终得到的平均loss都是一样的。

那pytorch贡献者也实现了这个loss求平均的功能,即通过gather的方式来求loss平均:

https://github.com/pytorch/pytorch/pull/7973/commits/c285b3626a7a4dcbbddfba1a6b217a64a3f3f3be

如果它们在一个有2个GPU的系统上运行,DP将采用多GPU路径,调用gather并返回一个向量。如果运行时有1个GPU可见,DP将采用顺序路径,完全忽略gather,因为这是不必要的,并返回一个标量。

4.关于内存

在ternimal下执行watch -n 1 nvidia-smi后会发现确实会使用多个GPU来并行训练。但是细心点会发现其实第一块卡的显存会占用的更多一些,那么这是什么原因导致的?

查阅pytorch官网的nn.DataParrallel相关资料,首先我们来看下其定义如下:

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

其中包含三个主要的参数:module,device_ids和output_device。

module即表示你定义的模型,device_ids表示你训练的device,output_device这个参数表示输出结果的device,而这最后一个参数output_device一般情况下是省略不写的,那么默认就是在device_ids[0],也就是第一块卡上,也就解释了为什么第一块卡的显存会占用的比其他卡要更多一些。进一步说也就是当你调用nn.DataParallel的时候,只是在你的input数据是并行的,但是你的output loss却不是这样的,每次都会在第一块GPU相加计算,这就造成了第一块GPU的负载远远大于剩余其他的显卡。

参考:

1.dataparallel:https://zhuanlan.zhihu.com/p/102697821

Guess you like

Origin blog.csdn.net/pearl8899/article/details/118640520