PyTorch判断是否使用GPU以及哪些数据需要放到GPU上

打印模型 model 和数据 data 在 gpu 上还是 cpu 上。

• 1. 通过判断模型 model 的参数是否在 cuda 上来判定模型是否在 gpu 上。

print ('Is model on gpu: ', next (model.parameters ()).is_cuda)
输出若是 True，则 model 在 gpu 上；若是 False，则 model 在 cpu 上。

• 2. 输出数据 data 的 device 字段。

print ('data device: ', data.device)
输出 gpu 则在 gpu 上，输出 cpu 则在 cpu 上。

哪些数据需要放到 GPU 上

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