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

Posted AI浩

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了.to(device)和.cuda()设置GPU的区别相关的知识,希望对你有一定的参考价值。

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

以上是关于.to(device)和.cuda()设置GPU的区别的主要内容,如果未能解决你的问题,请参考以下文章

PyTorch笔记: GPU上训练的模型加载到CPU/错误处理Attempting to deserialize object on a CUDA device but torch.cuda.is_a

TensorFlow报错:failed call to cuDevicePrimaryCtxRetain: CUDA_ERROR_INVALID_DEVICE

PyTorchcuda()与to(device)的区别

PyTorchcuda()与to(device)的区别

RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1.(

pytorch 多GPU 训练