PART 1: 使用GPU运行程序的方法
a 首先确定计算机或者服务器有可以使用的GPU
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()`
b 将网络转为cuda模式
net = Net().to(device)
注意:Net() 是已经定义好的网络类,用Net()来初始化net
c 将数据转为cuda模式
data.to(device)
data.cpu()
注意: 在使用GPU进行计算时,要搞清楚何时使用什么样式的数据,在计算时才不会出错
PART 2: CPU使用利用GPU训练好的模型
a 首先定义CPU计算
device = torch.device('cpu')
b 定义一个相同的在CPU情况下的网络
net = Net()
c 导入网络参数,并指定参数map_location=device,即网络的参数映射到基于CPU计算的网络
net.load_state_dict(torch.load("(已经训练好并且保存下来的模型地址).pkl", map_location=device))
下面就可以在CPU上直接使用GPU训练好的模型了!
PART 3: 几个小技巧
1)数据加载
data_loader = Data.DataLoader(dataset, batch_size=4, shuffle=True,
num_workers=8,pin_memory=True)
pin_memory=True,这里的参数将数据定义为锁页内存,这样转移到GPU上的速度会更快一点
2)需要什么在GPU上运行,就将什么内容传到GPU上。
.to(device)
PyTorch初学者,请多指教!!!
|