pytorch保存模型等相关参数，利用torch.save()，以及读取保存之后的文件
Posted qinduanyinghua
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了pytorch保存模型等相关参数，利用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)
其中dir表示保存文件的绝对路径+保存文件名，如‘/home/qinying/Desktop/modelpara.pth‘
二、
当你想恢复某一阶段的训练（或者进行测试）时，那么就可以读取之前保存的网络模型参数等。
checkpoint = torch.load(dir)
model.load_state_dict(checkpoint[‘net‘])
optimizer.load_state_dict(checkpoint[‘optimizer‘])
start_epoch = checkpoint[‘epoch‘] + 1
以上是关于pytorch保存模型等相关参数，利用torch.save()，以及读取保存之后的文件的主要内容，如果未能解决你的问题，请参考以下文章