![f9bd9d24dfae0ab1d62c858aa997e9cb.png](https://i-blog.csdnimg.cn/blog_migrate/1f6a0c810373e87b45ae55a27c9eb4f4.png)
遇到了
map_location=lambda storage, loc: storage
搞不懂,搜了一下,还是站在巨人的肩膀上好啊,哈哈哈哈哈!
pytorch在CPU和GPU上加载模型_人工智能_牧野的博客-CSDN博客blog.csdn.net![ea0002cd380efcade8bd1df901a5ecbe.png](https://i-blog.csdnimg.cn/blog_migrate/8c91bb69ba37f241ff05f4704f8643bd.png)
![a5e461818ae8e2be2b4129e46559ff49.png](https://i-blog.csdnimg.cn/blog_migrate/afb58b414fe7890c35d0b6655b85b5c7.png)
我们在使用pytorch的过程,经常会需要加载模型参数,不管是别人提供给我们的模型参数,还是我们自己训练的模型参数,那么加载模型参数就会碰到一些情况,即GPU模型和CPU模型,这两种模型是不能混为一谈的,下面分情况进行操作说明。
情况一:模型是GPU模型,预加载的训练参数也是GPU;模型是CPU模型,预加载的训练参数也是CPU,这种情况下我们都只用直接用下面的语句即可:
torch.load('model_dict.pkl')
情况二:GPU->CPU 模型是CPU,预加载的训练参数却是GPU,那么需要这样:
torch.load('model_dict.pkl', map_location=lambda storage, loc: storage)
情况三:CPU->GPU 模型是GPU,预加载的训练参数却是CPU:
torch.load('model_dic.pkl', map_location=lambda storage, loc: storage.cuda)
#CPU->GPU1 模型是GPU1,预加载的训练参数却是CPU:
torch.load('model_dic.pkl', map_location=lambda storage, loc: storage.cuda(1))