pytorch保存和加载训练模型torch.save()model.state_dict()torch.load()model.load_state_dictdel checkpoint

57 阅读 0 评论 38 点赞

我是靠谱客的博主仁爱画笔，最近开发中收集的这篇文章主要介绍pytorch保存和加载训练模型torch.save()model.state_dict()torch.load()model.load_state_dictdel checkpoint，觉得挺不错的，现在分享给大家，希望可以做个参考。

torch.save()

torch.save()可以保存一个字典，这里除了保存预训练模型的权重之外，还可以保存一些训练信息：
在这里插入图片描述

在这里插入图片描述

继承自torch.nn.Module的模型都具备的方法，可以获取模型的参数，返回的是字典：

module.state_dict().keys()
[‘bias’, ‘weight’]

配合torch.save（）可以保存模型。

torch.save(model.state_dict())
#或者
torch.save({'state_dict':model.state_dict()})

torch.load()可以加载一个字典：
在这里插入图片描述

继承自torch.nn.Module的模型都具备的方法：

checkpoint=torch.load('name.pth')
model.load_state_dict(checkpoint['state_dict'])#这里是之前保存的时候把模型参数存在了'state_dict'关键字下，这里看情况加载模型

加载到模型之后，删掉原有的checkpoint以节省内存。

del checkpoint  # dereference seems crucial
torch.cuda.empty_cache()#清掉之前的缓存

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

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