在使用model.load_state_dict进行权重文件加载时,类似keras中load_weights(by_name=True),是严格按照模型的中的各个模块的名称与权重文件中state_dict 的索引进行匹配的。如果不匹配,会报错 Error(s) in loading state_dict
。
注意
Pytorch与keras中load_weights(by_name=False)设置不同,由于Tensorflow是静态图
可以按照网络的拓扑结构,即模型与权重文件的的层次结构进行匹配
Pytorch是动态图
,即网络在训练的时候也可能会改变网络结构,因此当load_state_dict(, strict=False)时,出现不匹配的不会报错,但是对应的模块也无法加载字典索引对应的权重文件
故为实现权重文件的有效加载,需要对权重文件的字典索引按照我们重建的模型进行修改,这里我们使用到re模块,对索引字符串进行正则化匹配。
DenseNet
以torch官网给出的DenseNet修改方式为例。备注中有说明由于原权重文件中的网络层名为’norm.1’, ‘relu.1’, ‘conv.1’, ‘norm.2’, ‘relu.2’, ‘conv.2’.,而重建的网络中则直接使用norm1等方式,因此需要去除 ‘.’ 。
def _load_state_dict(model, model_url, progress):
# '.'s are no longer allowed in module names, but previous _DenseLayer
# has keys 'norm.1', 'relu.1', 'conv.1', 'norm.2', 'relu.2', 'conv.2'.
# They are also in the checkpoints in model_urls. This pattern is used
# to find such keys.
pattern = re.compile(
r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$')
state_dict = load_state_dict_from_url(model_url, progress=progress)
for key in list(state_dict.keys()):
res = pattern.match(key)
if res:
new_key = res.group(1) + res.group(2)
state_dict[new_key] = state_dict[key]
del state_dict[key]
model.load_state_dict(state_dict)
这里的关键是操作便是,匹配的正则表达式
pattern = re.compile(r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$')
权重文件中的索引
‘features.denseblock1.denselayer1.norm.1.weight’
模型中的索引
‘features.denseblock1.denselayer1.norm1.weight’
第一次看可能有点懵,确实写的有点复杂,我们从左到右依次来看:
r’ 后面接需要匹配的内容
^ 匹配字符串的开始;
. 转义字符匹配.
*重复零次或更多次
\d 匹配数字
?:norm|relu|conv 匹配三组字符串中的一类
?:[12] 匹配数字1或2
$ 匹配字符串的结尾
其中res.group(*)应该是根据r’后面的括号进行划分的不同group
res.group(0)
‘features.denseblock1.denselayer1.norm.1.weight’
res.group(1)
‘features.denseblock1.denselayer1.norm’
res.group(2)
‘1.weight’
re模块参考
https://www.runoob.com/python/python-reg-expressions.html
https://www.cnblogs.com/shenjianping/p/11647473.html