pytorch CPU/GPU模型的保存和加载

最新推荐文章于 2024-07-31 22:40:07 发布

Duwr

最新推荐文章于 2024-07-31 22:40:07 发布

阅读量756

点赞数

文章标签：深度学习

本文链接：https://blog.csdn.net/weixin_42094192/article/details/107135945

版权

情况一：模型是GPU模型，预加载的训练参数也是GPU；模型是CPU模型，预加载的训练参数也是CPU，这种情况下我们都只用直接用下面的语句即可：
torch.load(‘model_dict.pkl’)
情况二：模型是CPU，预加载的训练参数却是GPU，那么需要这样：
torch.load(‘model_dict.pkl’, map_location=lambda storage, loc: storage)
情况三：模型是GPU，预加载的训练参数却是CPU：
torch.load(‘model_dic.pkl’, map_location=lambda storage, loc: storage.cuda)