Pytorch——将模型load到gpu或cpu上

74 阅读 0 评论 49 点赞

我是靠谱客的博主友好金针菇，最近开发中收集的这篇文章主要介绍Pytorch——将模型load到gpu或cpu上，觉得挺不错的，现在分享给大家，希望可以做个参考。

很多时候我们在gpu上训练一个模型，但是在inference的时候不想使用gpu。或者想在别的gpu上使用，那么怎么办呢？

需要在load的时候就选择device。

保存了模型的参数(model.state_dict())到文件model.pth中。

1、cpu->cpu 或gpu->gpu

这种情况是最简单的：

checkpoint = torch.load('model.pth')

model.load_state_dict(checkpoint)

2、cpu->gpu1

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

3、gpu 0 -> gpu 1

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

4、gpu -> cpu

checkpoint=torch.load('model.pth', map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint)

以上就是友好金针菇为你收集整理的Pytorch——将模型load到gpu或cpu上的全部内容，希望文章能够帮你解决Pytorch——将模型load到gpu或cpu上所遇到的程序开发问题。

如果觉得靠谱客网站的内容还不错，欢迎将靠谱客网站推荐给程序员好友。

本图文内容来源于网友提供，作为学习参考使用，或来自网络收集整理，版权属于原作者所有。