赞
踩
1. 直接使用CPU训练时,将.cuda改成.cpu即可。或(1条消息) 将gpu运行的torch程序改为cpu运行_油泼西红柿的博客-CSDN博客_torch使用cpuhttps://blog.csdn.net/qq_37813206/article/details/113132421
2. 使用GPU训练好的模型,加载在CPU上测试。
首先,将代码中的.cuda都改成.cpu。然后,因为在GPU上训练的模型参数比CPU训练的参数多了“modlue.”,所以加载时需要去掉。
- # ------------GPU-----------------
- # os.environ["CUDA_VISIBLE_DEVICES"] = '0,1' # 指定要使用的显卡
- # model = torch.nn.DataParallel(model) # 使用多个显卡运行模型
- # model.load_state_dict(torch.load('xxx_50.pkl')) # 加载模型
- # model = model.cuda() # 将模型加载到 GPU
- # ------------CPU-----------------
- state_dict = torch.load('xxx_50.pkl',map_location=torch.device('cpu'))
- state_dict_new = OrderedDict()
- for k, v in state_dict.items():
- name = k[7:] # 去掉 "module."
- state_dict_new[name] = v
- model.load_state_dict(state_dict_new) # 加载模型
- # -------------END----------------
- model.eval()
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。