保存模型的几种格式
.pt .pth .pkl:它们并不是在格式上有区别,只是后缀不同而已(仅此而已)
- 方式一:只保存模型权重参数,不保存模型结构
torch.save(model.state_dict(), 'mymodel.pth') # 只保存模型权重参数,不保存模型结构
model = My_model(*args, **kwargs) # 这里需要重新模型结构,My_model
model.load_state_dict(torch.load('mymodel.pth')) # 这里根据模型结构,调用存储的模型参数
model.eval()
举例
dir = 'mymodel.pth'
state = {'net':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}
torch.save(state, dir)
checkpoint = torch.load(dir) # checkpoint 把之前save的state加载进来
model.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch = checkpoint['epoch'] + 1
torch.save(model, 'mymodel.pth') # 保存整个model的状态
model=torch.load('mymodel.pth') # 这里已经不需要重构模型结构了,直接load就可以
model.eval()
资料源自网络,未标注来源,仅作为自我QA记录
————————————————
版权声明:本文为CSDN博主「木槿qwer」的原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/lwqian102112/article/details/121553561