1、目前主流方法:.to(device)方法(推荐)
import torch
import time
#1.通常用法
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
data = data.to(device)
model = model.to(device)
'''
1.先创建device对象
2.to函数指定数据或者模型放到哪里
'''
#2.将构建的tensor或者模型放到指定设备上(GPU)
torch.device('cuda',0) #这里的0指的是设备的序号
torch.device('cuda:0')
#3.例子 cpu转到GPU
s = time.time()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
a = torch.rand([1024, 1024, 10]).to(device)
print('time:{:6.3f}'.format(time.time()-s)) # 输出: time: 0.087
#4.例子2 Tensor GPU 转到cpu运行
predict_boxes = predictions["boxes"].to("cpu").numpy()
predict_classes = predictions["labels"].to("cpu").numpy()
predict_scores = predictions["scores"].to("cpu").numpy()
.cuda()方法和.to(device)方法耗时基本差不多。但是 .to(device)代码风格更容易后期修改。
2、to(device)需要注意的是:
使用GPU训练的时候,需要将Module对象和Tensor类型的数据送入到device。通常会使用 to.(device)。但是需要注意的是:
- 对于Tensor类型的数据,使用to.(device) 之后,需要接收返回值,返回值才是正确设置了device的Tensor。
- 对于Module对象,只用调用to.(device) 就可以将模型设置为指定的device。不必接收返回值。
Module对象设置device的写法
model.to(device)
Tensor类型的数据设置 device 的写法。
samples = samples.to(device)
当然model也可以接收返回值。
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
img = img.to(device)
官网解释:
3、如果你有多个GPU
这样可以显示指定需要使用的计算资源,特别是有多个GPU的情况下。
那么可以参考以下代码:
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = Model()
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model,device_ids=[0,1,2])
model.to(device)
1、torch.device用法总结_会飞的萨摩的博客-CSDN博客主要参考这两篇文章:
1、torch.device用法总结_会飞的萨摩的博客-CSDN博客
2、pytorch并行处理详解(多GPU,环境变量)_xys430381_1的博客-CSDN博客
Original: https://blog.csdn.net/qimo601/article/details/123822178
Author: 三世
Title: python指定Gpu
原创文章受到原创版权保护。转载请注明出处:https://www.johngo689.com/672855/
转载文章受原作者版权保护。转载请注明原作者出处!