pytorch 多gpu并行训练

单机多卡并行训练

torch.nn.DataParallel

我一般在使用多GPU的时候, 会喜欢使用os.environ['CUDA_VISIBLE_DEVICES']来限制使用的GPU个数, 例如我要使用第0和第3编号的GPU, 那么只需要在程序中设置：

os.environ['CUDA_VISIBLE_DEVICES'] = '0,3'

这个参数最好在主体代码开始的时候写上，因为它要写在model调用cuda函数前使用。

例如我们需要训练的文件为train.py，也是用‘0,3’号gpu，可以在代码外部的命令行中输入命令：

CUDA_VISIBLE_DEVICES='0,3' python train.py

这样的方式有两点好处，1.不用再关心代码是否在调用cuda前已经设置好参数；2.写在外面方便调用空闲gpu，不用每次使用的时候再去代码中修改。

如果是模型, 那么需要执行下面的这几句代码：

model = nn.DataParallel(model)
model = model.cuda()

如果是数据, 那么直接执行下面这几句代码就可以了：

inputs = inputs.cuda()
labels = labels.cuda()

 

 Reference:

https://zhuanlan.zhihu.com/p/86441879