深度模型训练时CPU或GPU的使用model.to(device)

netouch 2024-05-16 发布于北京

展开全文

一、使用device控制使用CPU还是GPU

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 
# 单GPU或者CPU.先判断机器上是否存在GPU，没有则使用CPU训练
model = model.to(device)
data = data.to(device)
#或者在确定有GPU的情况下，直接使用
model = model.cuda()
data = data.cuda()
#在确定没有GPU的情况下，直接使用
model = model.cpu()
data = data.cpu()

注意：

1.tensor和numpy都是矩阵，前者能在GPU上运行，后者只能在CPU运行，所以要注意数据类型的转换。

2.当你使用model.to(device)时，它会将模型的参数和缓冲区移动到指定的设备上。而当你使用model = model.to(device)时，它会将整个模型移动到指定的设备上。

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

建议使用model.to(device)的方式，这样可以显示指定需要使用的计算资源，特别是有多个GPU的情况下，可以并行处理，加快速度。

参考.to(device)和.cuda()设置GPU的区别_.cuda()和to(device)-CSDN博客

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

三、指定使用的GPU

使用方式

import os
# 给服务器上的GPU编号
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'
# 指定要使用哪张卡
device_ids = [0, 1] #假设只用两张卡
# 将模型搬到GPU上，并行化处理
model = torch.nn.DataParallel(model, device_ids=device_ids)
model = model.cuda()
# 或者直接
model = torch.nn.DataParallel(model, device_ids=device_ids).cuda()