在使用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模块,对索引字符串进行正则化匹配。
以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’
https://www.runoob.com/python/python-reg-expressions.html https://www.cnblogs.com/shenjianping/p/11647473.html
