1. Pytorch 加载完整模型的参数
- 保存加载整个模型
# 保存整个模型
torch.save (model_object, 'model.pk1')
# 加载整个模型
model = torch.load('model.pkl')
- 保存模型的参数 (推荐使用)
# 模型参数保存
torch.save (model_object.state_dict(), 'params.pk1') # 保存的是以字典 key - value pair 形式的数据,每一个参数对应着一个值 state_dict 状态字典
# 模型参数加载
model_object.load_state_dict(torch.load('params.pkl'))
2.加载部分模型参数
大多数时候我们需要根据我们的任务调节我们的模型,很难保证说我们的模型和我们的任务完全一致,但是预训练模型确实有利于提高任务的准确率,为了结合两者的优点,我们可以使用加载部分预训练模型的方法。
# 0. 先加载进全部的参数存进 pretrained_dict 中
pretrained_dict = t.load ('params.pk1')
#1. load state dict 加载状态字典,就是加载属于这个 model 的状态字典,看看这个 model 里面有哪些变量的 name
model_dict = model.state_dict()
# 2. filter out unnecessary keys 过滤掉不属于这个模型的 key 值,只有当 key 的值在这个模型里面的时候才加载这个 key 其他的 key 值自动放弃使用 相当于两个模型求了一个交集。
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 3. overwrite entries in the existing state dict 更新 model 的数据字典 (更新成为可以加载的数据字典)
model_dict.update(pretrained_dict)
# 4. load the new state dict () 加载参数,只加载模型中存在的参数,放弃一些不存在的参数,只加载修改过的参数。
model.load_state_dict(model_dict)