由 torch.save(model.state_dict(), checkpoint_path) 展开的学习

资料仅作为自己复习回顾
参考资料1

torch.save功能
保存模型等相关参数,以及读取保存之后的文件

一、如何保存模型参数、优化器参数等

model = Net() # 网络
optimizer = optim.Adam(model.parameters(), lr=args.lr) # 优化器
# 假设在某个epoch,我们要保存模型参数,优化器参数以及epoch
# 1、先建立一个字典,保存三个参数
state = {‘net’:model.state_dict(), ‘optimizer’:optimizer.state_dict(), ‘epoch’:epoch}
# 2、调用torch.save() 保存模型
torch.save(state, dir)

可以看到,我们的state可以只有模型参数,也可以包含优化器、epoch参数

二、对保存的模型我们要如何读取

当你想恢复某一阶段的训练(或者进行测试)时,那么就可以读取之前保存的网络模型参数等。

checkpoint = torch.load(dir)
model.load_state_dict(checkpoint[‘net’])
optimizer.load_state_dict(checkpoint[‘optimizer’])
start_epoch = checkpoint[‘epoch’] + 1

当前盲区:
model.load_state_dict
optimizer.load_state_dict

你可能感兴趣的:(pytorch,深度学习,python)