pytorch保存模型的两种方式
一、直接保存整个模型并读取## 保存模型torch.save(your_model, 'model_name.pth')## 读取模型your_model = torch.load('model_name.pth')二、只保存模型中的参数并读取## 保存模型torch.save({'you_model': your_model.state_dict()}, 'model_name.pt')## 读
文章共482字 · 阅读需要大约2分钟
一键AI生成摘要,助你高效阅读
问答
·
一、直接保存整个模型并读取
# 创建你的模型实例对象: model
model = net()
## 保存模型
torch.save(model, 'model_name.pth')
## 读取模型
model = torch.load('model_name.pth')
二、只保存模型中的参数并读取
## 保存模型
torch.save({'model': model.state_dict()}, 'model_name.pth')
## 读取模型
model = net()
state_dict = torch.load('model_name.pth')
model.load_state_dict(state_dict['model'])
总结
我们可以看到第一种方法可以直接保存模型,加载模型的时候直接把读取的模型给一个参数就行。而第二种方法则只是保存参数,在读取模型参数前要先定义一个模型(模型必须与原模型相同的构造),然后对这个模型导入参数。虽然麻烦,但是可以同时保存多个模型的参数,而第一种方法则不能,而且第一种方法有时不能保证模型的相同性(你读取的模型并不是你想要的)。
总的来说,我们一般来选择第二种来保存和读取。
退一步讲,如何保存模型决定了如何读取模型。
更多推荐
已为社区贡献1条内容
所有评论(0)