.to(device)和.cuda()设置GPU的区别

AI浩

已于 2022-05-11 21:41:04 修改

阅读量9.1k

点赞数 14

分类专栏：人工智能文章标签： cuda gpu 深度学习 pytorch

于 2021-07-09 11:29:59 首次发布

本文为博主原创文章，未经博主允许不得转载。

本文链接：https://blog.csdn.net/hhhhhhhhhhwwwwwwwwww/article/details/118602874

版权

人工智能专栏收录该内容

178 篇文章 96 订阅

订阅专栏

.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_DEVICES']="1"
model.cuda()
#如果是多GPU
os.environ['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()

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