本文最后更新于238 天前,其中的信息可能已经过时,如有错误请发送邮件到1729915388@qq.com
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)
os.environ[‘CUDA_VISIBLE_DEVICES’] = ‘1’
os.environ[‘CUDA_VISIBLE_DEVICES’] = ‘1,2’