最近在尝试用torch单机多卡进行训练。 网上有很多方法,有的讲的也很详细,但是torch版本更新的还是很快的。所以自己也踩了很多坑。在这里记录下来,希望对大家有帮助。
本文适用torch版本:1.10
torch单机多gpu训练有两种方式
torch.nn.DataParallel
torch.nn.DataParallel(
module,
device_ids=None,
output_device=None,
dim=0)
这一种方式是比较老的,官方现在也不太推荐,相对来说,这种方式的速度并不算太快。
但是使用起来还算简单,需要修改的代码并不多。
在训练RNN的时候,坑非常多。可以看看博客:RNN与torch DataParallel的爱恨情仇
下面介绍一下相对于普通的训练,这种方式需要进行哪些修改。
使用方式:
from torch.nn import DataParallel
...
model = nn.DataParallel(mode