1. 利用CUDA_VISIBLE_DEVICES设置可用显卡

放在网络net开始训练的代码之前

os.environ['CUDA_VISIBLE_DEVICES'] = gpu_ids 

需要放在import torch之前？否则没有效果。

该行代码的意思是只选择指定标号GPU进行使用，执行完该行代码之后python环境无法检测到指定GPU之外的其他GPU

2 .cuda()方法和torch.cuda.set_device()

还可以使用.cuda()[包括model.cuda()/loss.cuda()/tensor.cuda()]方法和torch.cuda.set_device()来把模型和数据加载到对应的gpu上。

(1) .cuda()

以model.cuda()为例，加载方法为：

model.cuda(gpu_id) # gpu_id为int类型变量，只能指定一张显卡
model.cuda('cuda:'+str(gpu_ids)) #输入参数为str类型，可指定多张显卡
model.cuda('cuda:1,2') #指定多张显卡的一个示例

(2) torch.cuda.set_device()

使用torch.cuda.set_device()可以更方便地将模型和数据加载到对应GPU上, 直接定义模型之前加入一行代码即可

torch.cuda.set_device(gpu_id) #单卡
torch.cuda.set_device('cuda:'+str(gpu_ids)) #可指定多卡

多卡数据并行torch.nn.DataParallel

多卡数据并行一般使用

torch.nn.DataParallel(model,device_ids)
# self.net = torch.nn.DataParallel(self.net, device_ids=list(range(torch.cuda.device_count())))

其中model是需要运行的模型，device_ids指定部署模型的显卡，数据类型是list

device_ids中的第一个GPU（即device_ids[0]）和model.cuda()或torch.cuda.set_device()中的第一个GPU序号应保持一致，否则会报错。此外如果两者的第一个GPU序号都不是0,比如设置为：

model=torch.nn.DataParallel(model,device_ids=[2,3])
model.cuda(2)

那么程序可以在GPU2和GPU3上正常运行，但是还会占用GPU0的一部分显存（大约500M左右），这是由于pytorch本身的bug导致的（截止1.4.0，没有修复这个bug）。

device_ids的默认值是使用可见的GPU，不设置model.cuda()或torch.cuda.set_device()等效于设置了model.cuda(0)