我是靠谱客的博主 仁爱画笔，这篇文章主要介绍pytorch保存和加载训练模型torch.save()model.state_dict()torch.load()model.load_state_dictdel checkpoint，现在分享给大家，希望可以做个参考。
pytorch保存和加载训练模型
- torch.save()
- model.state_dict()
- torch.load()
- model.load_state_dict
- del checkpoint
torch.save()
torch.save()可以保存一个字典，这里除了保存预训练模型的权重之外，还可以保存一些训练信息：
model.state_dict()
继承自torch.nn.Module的模型都具备的方法，可以获取模型的参数，返回的是字典：
module.state_dict().keys()
[‘bias’, ‘weight’]
配合torch.save（）可以保存模型。
复制代码
1
2
3
4torch.save(model.state_dict()) #或者 torch.save({'state_dict':model.state_dict()})
torch.load()
torch.load()可以加载一个字典：
model.load_state_dict
继承自torch.nn.Module的模型都具备的方法：
复制代码
1
2
3checkpoint=torch.load('name.pth') model.load_state_dict(checkpoint['state_dict'])#这里是之前保存的时候把模型参数存在了'state_dict'关键字下，这里看情况加载模型
del checkpoint
加载到模型之后，删掉原有的checkpoint以节省内存。
复制代码
1
2
3del checkpoint # dereference seems crucial torch.cuda.empty_cache()#清掉之前的缓存
最后
以上就是仁爱画笔最近收集整理的关于pytorch保存和加载训练模型torch.save()model.state_dict()torch.load()model.load_state_dictdel checkpoint的全部内容，更多相关pytorch保存和加载训练模型torch.save()model.state_dict()torch.load()model.load_state_dictdel内容请搜索靠谱客的其他文章。
本图文内容来源于网友提供，作为学习参考使用，或来自网络收集整理，版权属于原作者所有。
发表评论 取消回复