pytorch如何将用GPU训练好的网络模型用到CPU中

本文参考链接地址:
1.https://blog.csdn.net/bc521bc/article/details/85623515
2.https://www.ptorch.com/news/74.html

刚开始用python 小小小菜鸟一枚,啥也不会,在GitHub上下载了代码,运行后错误百出,陷入了深深的绝望之中。无奈只好一点点的差,一点点的改。

这篇文章简单记一下:如何将用GPU训练好的网络模型用到CPU中
源代码如下:

        feature_encoder.load_state_dict(torch.load(str("./models/miniimagenet_feature_encoder_" + str(CLASS_NUM) +"way_" + str(SAMPLE_NUM_PER_CLASS) +"shot.pkl")))       

在没有GPU的电脑上运行,报错如下:

        RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location='cpu' to map your storages to the CPU.       

我做的更改如下:

        feature_encoder.load_state_dict(torch.load(str("./models/miniimagenet_feature_encoder_" + str(CLASS_NUM) +"way_" + str(SAMPLE_NUM_PER_CLASS) +"shot.pkl"),map_location=lambda storage, loc: storage))

改完后,目前还没发现其他错误。
主要是因为我也还没有运行到最后,我还得一点点再看看,希望我改对了[捂脸笑]

最后,记录下cpu和GPU间相互转换的方法:

gpu→ cpu

torch.load('parameters.pth', map_location=lambda storage, loc: storage)

cpu—>gpu 1

torch.load('parameters.pth', map_location=lambda storage, loc: storage.cuda(1))

gpu1—>gpu0

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

最后,同类之间,也就是cpu—>cpu或者gpu—>gpu:

checkpoint = torch.load(‘parameters.pth’)

不对的地方,还请各位大神多多指教哦!!!

  • 1
    点赞
  • 7
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值