赞
踩
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device
注意:
a) 只有一块GPU则使用:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
b) 有多块GPU则使用:
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
cuda:1----定义使用GPU1,GPU的数量通常从0开始排序
c) 获取能够使用的GPU数量
`torch.cuda.device_count()`
net = Net().to(device)
注意:Net() 是已经定义好的网络类,用Net()来初始化net
data.to(device) #将数据转化为使用指定的GPU可计算形式0
data.cpu() #将GPU可计算的形式的数据转化为CPU的可计算形式
注意:
在使用GPU进行计算时,要搞清楚何时使用什么样式的数据,在计算时才不会出错
device = torch.device('cpu')
net = Net()
net.load_state_dict(torch.load("(已经训练好并且保存下来的模型地址).pkl", map_location=device))
下面就可以在CPU上直接使用GPU训练好的模型了!
data_loader = Data.DataLoader(dataset, batch_size=4, shuffle=True,
num_workers=8,pin_memory=True)
pin_memory=True,这里的参数将数据定义为锁页内存,这样转移到GPU上的速度会更快一点
.to(device)
PyTorch初学者,请多指教!!!
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。