码迷,mamicode.com
首页 > 其他好文 > 详细

【colab pytorch】保存模型

时间:2020-03-09 21:11:08      阅读:200      评论:0      收藏:0      [点我收藏+]

标签:check   span   mode   poi   col   训练   star   save   optimizer   

保存模型总体来说有两种:

第一种:保存训练的模型,之后我们可以继续训练

(1)保存模型

state = { model: model.state_dict(), optimizer:optimizer.state_dict(), epoch: epoch }   
torch.save(state, path)

model.state_dict():模型参数

optimizer.state_dict():优化器

epoch:保存epoch,为了可以接着训练

(2)恢复模型

checkpoint = torch.load(path)
model.load_state_dict(checkpoint[model])
optimizer.load_state_dict(checkpoint[optimizer])
start_epoch = checkpoint[epoch]+1

第二种:保存测试的模型,一般保存准确率最高的

(1)保存模型

这时我们只需要保存模型参数就行了

torch.save(model.state_dict, path)

(2)恢复模型

model.load_state_dict(torch.load(path))

 

【colab pytorch】保存模型

标签:check   span   mode   poi   col   训练   star   save   optimizer   

原文地址:https://www.cnblogs.com/xiximayou/p/12450761.html

(0)
(0)
   
举报
评论 一句话评论(0
登录后才能评论!
© 2014 mamicode.com 版权所有  联系我们:gaon5@hotmail.com
迷上了代码!