pytorch并行

首先指定GPU：

两种方法：1.在命令行指定：CUDA_VISIBLE_DEVICES=3,0,1,2 python main.py 

　　　　　　也就是在普通的python main.py 命令前加 CUDA_VISIBLE_DEVICES=gpu编号

　　　　　　注意：等号两边不能有多余的空格

　　　　　2.在main.py最前面加入os.environ['CUDA_VISIBLE_DEVICES'] = '1,3'（当然是在import os之后）

指定的GPU的顺序第一个是主GPU，显存使用会比其他的大一点

 

然后再单GPU基础上改一行代码：

model = model.cuda()改为

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