【colab pytorch】保存模型

2020-08-26 10:18:40 浏览数 (2)

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

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

(1)保存模型

代码语言:javascript复制
state = { 'model': model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch': epoch }   
torch.save(state, path)

model.state_dict():模型参数

optimizer.state_dict():优化器

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

(2)恢复模型

代码语言:javascript复制
checkpoint = torch.load(path)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch = checkpoint['epoch'] 1

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

(1)保存模型

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

代码语言:javascript复制
torch.save(model.state_dict, path)

(2)恢复模型

代码语言:javascript复制
model.load_state_dict(torch.load(path))

0 人点赞