pytorch Error(s) in loading state_dict 使用re模块 改变字典索引

在使用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