os.environ[‘CUDA_VISIBLE_DEVICES‘]使用问题记录

最新推荐文章于 2024-09-08 18:03:17 发布

又又酸又菜又多余

最新推荐文章于 2024-09-08 18:03:17 发布

阅读量2.9k

点赞数 2

分类专栏：问题记录文章标签：深度学习 python pytorch

本文链接：https://blog.csdn.net/qq_43470698/article/details/126919963

版权

问题记录专栏收录该内容

3 篇文章 0 订阅

订阅专栏

os.environ['CUDA_VISIBLE_DEVICES'] = '0，1(卡号)'

1、位置

至于import torch之前

2、和device映射关系

os.environ['CUDA_VISIBLE_DEVICES'] = '6,1,2'
device = torch.device("cuda:0")

这里device的cuda0指向的是卡6，cuda2指向卡2

3、.cuda()和.to(device）区别

摘自pytorh .to(device) 和.cuda()的区别
.to(device) 可以指定CPU 或者GPU

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # 单GPU或者CPU
model.to(device)
#如果是多GPU
if torch.cuda.device_count() > 1:
  model = nn.DataParallel(model，device_ids=[0,1,2])
model.to(device)

.cuda()只能指定GPU

#指定某个GPU
os.environ['CUDA_VISIBLE_DEVICE']='1'
model.cuda()
#如果是多GPU
os.environment['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
device_ids = [0,1,2,3]
net  = torch.nn.Dataparallel(net, device_ids =device_ids)
net  = torch.nn.Dataparallel(net) # 默认使用所有的device_ids 
net = net.cuda()