torch gpu改cpu

一、

torch在训练前一般会加以下代码：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

有cuda[gpu] 就使用，没有就用cpu

之后代码中进行修改：

model=xxx.cuda

一律改为

model=xxx.to（device） 

这种方法：如果电脑有gpu可用，调用的还是gpu

二、

(1) 模型在GPU上保存，运行在CPU上

torch.save(model.state_dict(), PATH)

device = torch.device("cpu")
model = xxxxxxx(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))
or

torch.load(xxxxxxxxxxxx) 同上

model.load_state_dict(没有map_location 参数）会报错

（2）模型在cpu,运行在gpu：

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))
model.to(device)

如果想控制使用哪个gpu：

torch.load('modelparameters.pth', map_location={'cuda:1':'cuda:0'}) 

---

---

---

ps：

import torch
print(torch.cuda.device_count())  # 可用gpu数量
print(torch.cuda.is_available())  # 是否可用gpu