colab pytorch保存模型
Posted xiximayou
tags:
篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了colab pytorch保存模型相关的知识,希望对你有一定的参考价值。
保存模型总体来说有两种:
第一种:保存训练的模型,之后我们可以继续训练
(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保存模型的主要内容,如果未能解决你的问题,请参考以下文章
如何将自定义 Pytorch 模型转换为 torchscript(pth 到 pt 模型)?