pytorch保存模型非常简单,主要有两种方法:
只保存参数;(官方推荐)
保存整个模型 (结构+参数)。
由于保存整个模型将耗费大量的存储,故官方推荐只保存参数,然后在建好模型的基础上加载。本文介绍两种方法,但只就第一种方法进行举例详解。
一、只保存参数
1.保存
torch.save(model.state_dict(), path)
特别地,如果还想保存某一次训练采用的优化器、epochs等信息,可将这些信息组合起来构成一个字典,然后将字典保存起来:
state = {'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch}
torch.save(state, path)
2.加载
针对上述第一种情况,也只需要一句即可加载模型:
model.load_state_dict(torch.load(path))
针对上述第二种以字典形式保存的方法,加载方式如下:
checkpoint = torch.load(path)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
epoch = checkpoint(['epoch'])
在加载的模型基础上继续训练
def main():
# 如果test_flag=True,则加载已保存的模型
if test_flag:
# 加载保存的模型直接进行测试机验证,不进行此模块以后的步骤
checkpoint = torch.load(log_dir)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch = checkpoint['epoch']
test(model, test_load)
return
# 如果有保存的模型,则加载模型,并在其基础上继续训练
if os.path.exists(log_dir):
checkpoint = torch.load(log_dir)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch = checkpoint['epoch']
print('加载 epoch {} 成功!'.format(start_epoch))
else:
start_epoch = 0
print('无保存模型,将从头开始训练!')
for epoch in range(start_epoch+1, epochs):
train(model, train_load, epoch)
test(model, test_load)
# 保存模型
state = {'model':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}
torch.save(state, log_dir)