为什么要用dataparallel?
加速训练?原理是什么?
最根本的原因是加大了batch_size,所以加速了训练。原因是我们在用GPU训练模型的时候,因为模型本身占现存就很大了,比方说BERT,那么在一次forward的过程中,如果batch_size过大,那么单块板子上计算图就构建地太大了,因此有可能炸显存,所以也就限制了batch_size的大小。
而dataparallel的原理就是把原先的model replicate多份,每份都是一个副本,然后在forward的时候,就可以一次性传入更大的batch_size。
比方说,原先模型在单卡上只能最多feed 16 batch_size的样本,现存就上限了;现在用dataparallel把model复制到3块板子上,那就可以forward 16*3 也就是48 batch_size的样本,这些样本会被均分为三份,分别给三块板子上的那三个parallel model(如下图);
我们都知道,一个batch 的loss其实就是把batch里面所有单个样本的loss取avg,所以我们把一个batch 48分为更小的三份,每份16,最后在三块板子上的loss再取一次avg,这和直接forward 48 batch_size其实相差不大。
虽然也有人做过实验,发现均分到多块板子上,每块板子算出来的loss其实会有比较大的差距的,一般如果你的实验受batch_size影响比较大(比方说metric learning),那么就最好别用dataparallel。
其他tips:
-
被封装成dataparallel的model,forward的时候最后就返回loss,尽量不要返回额外的feature、prediction,这样会导致主卡和其他几块副卡负载不均衡。还可以通过其他方法缓解这个问题:
和nn.DataParallel说再见 -
有些情况下,如果你输入的batch不是n_gpu的倍数,那么会报错;同理,你的forward如果返回的类型不是cuda上的tensor,或者tensor的除batch_size以外的维度不一致,那么也会报错(多卡的结果不能gather到一起)。所以尽量推荐在eval的时候不要用dataparallel去预测,因为有时候我们可能想在test的时候获得模型的预测结果,而预测的返回内容不一定会是tensor,而且sample的batch_size也不能完全做到是n_gpu的倍数。
可以用:
在pred或者eval的时候还原回单卡model -
保存的时候注意保存主model: