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
 

文章出处登录后可见!

已经登录?立即刷新

共计人评分,平均

到目前为止还没有投票!成为第一位评论此文章。

(0)
xiaoxingxing的头像xiaoxingxing管理团队
上一篇 2023年2月25日 下午8:15
下一篇 2023年2月25日 下午8:16

相关推荐