【colab pytorch】保存模型

时间:2022-07-23
本文章向大家介绍【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))

随机文章
本站知识点必读
最近更新