torch.nn.DataParallel的一些tip(持续更新)

为什么要用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：