PyTorch用预训练模型初始化网络的部分权重

PyTorch用预训练模型初始化网络的部分权重

关于从预训练的模型加载参数时可能会遇到的一些阻碍,笔者对自己的绵薄经验在此进行分享,也便于自己之后的温习。

加载初始化参数可参照论坛中 apaszke的做法,其删除了与当前model不匹配的key。
代码为:

pretrained_dict = ...
model_dict = model.state_dict()
# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
# 3. load the new state dict
model.load_state_dict(model_dict)

如果修改后的模型与预训练模型中无key重名现象,此法可以使用。但是如果有重名现象,则可能出现训练出错,如加载参数与模型的正确维度不匹配等。此时,加一个尺寸相等的限制条件即可。代码如下:

pretrained_dict = ...
model_dict = model.state_dict()

pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict 
   					and v.size() == model_dict[k].size()}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)

此方法参考了论坛中 ImgPrcSng的做法。

你可能感兴趣的:(笔记,pytorch,深度学习)