Pytorch模型保存

Pytorch模型保存

.pth    .pkl    .chkpl 是python的存储文件

.pth文件详解

在pytorch进行模型保存的时候，一般有两种保存方式，一种是保存整个模型，另一种是只保存模型的参数。

torch.save(model.state_dict(), "my_model.pth")  # 只保存模型的参数
 
torch.save(model, "my_model.pth")  # 保存整个模型

保存的模型参数实际上一个字典类型，通过key-value的形式来存储模型的所有参数，本文以自己在实践过程中使用的一个.pth文件为例来说明，使用的是整个模型。

模型四个键值分别详解

1. net["model"] 详解    键model所对应的值是一个OrderedDict，而这个OrderedDict字典里面又存储着所有的每一层的参数名称以及对应的参数值。
2. net["optimizer"]详解
3. net["scheduler"] 详解
4. net["iteration"] 详解

.pkl文件详解

#只保存模型参数
# 保存
torch.save(model.state_dict(), '\parameter.pkl')
# 加载
model = TheModelClass(...)
model.load_state_dict(torch.load('\parameter.pkl'))

#保存完整模型
# 保存
torch.save(model, '\model.pkl')
# 加载
model = torch.load('\model.pkl')