深度学习保存模型——.pt .pth .pkl

保存模型的几种格式
.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

• 方式二：保存整个model的状态

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