【colab pytorch】保存模型
【colab pytorch】保存模型
西西嘛呦
发布于 2020-08-26 10:18:40
发布于 2020-08-26 10:18:40
代码可运行
运行总次数：0
代码可运行
保存模型总体来说有两种：
第一种：保存训练的模型，之后我们可以继续训练
（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))本文参与 腾讯云自媒体同步曝光计划，分享自作者个人站点/博客。
原始发表：2020-03-09 ，如有侵权请联系 cloudcommunity@tencent.com 删除
评论
登录后参与评论
推荐阅读
腾讯云开发者
