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()



原文地址:https://www.cnblogs.com/jiangnanyanyuchen/p/13603072.html