pytorch 加载预训练模型

pytorch 提供了 state_dict() 和 load_state_dict() 两个参数用来保存和加载模型参数, 前者将模型参数保存为字典形式, 后者将字典形式的模型参数载入到模型当中. 

但是, 使用 load_state_dict() 加载模型参数时, 要求保存的模型参数键值类型和模型完全一致, 一旦我们对模型结构做了些许修改, 就不能直接调用该函数.

这时候, 可以采取用保存的模型参数来更新当前模型参数的方式, 达到加载预训练模型的目的.

1. 首先, 读取当前模型参数

model_dict = model.state_dict()

2. 读取预训练模型, 并选取要保留的部分

pre_dict = torch.load('path')
pre_dict = {k: v for k, v in pre_dict.items() if k in model_dict}

3. 使用预训练的模型更新当前模型参数

model_dict.update(pre_dict)

4. 加载模型参数

model.load_state_dict(model_dict)