Pytorch加载预训练模型，修改变动层参数(如类别数)

男娘i 2023-01-17 08:49 198阅读 0赞

问题：加载预训练模型时若出现如下报错，则未考虑预训练模型与当前网络结构不相同层，如最终类别数与公共数据训练的模型类别数不一致。

RuntimeError: Error(s) in loading state_dict for DistributedDataParallel:
    	size mismatch for module.rpn.head.cls_logits.weight: copying a param with shape torch.Size([80, 256, 3, 3]) from checkpoint, the shape in current model is torch.Size([27, 256, 3, 3]).
    	size mismatch for module.rpn.head.cls_logits.bias: copying a param with shape torch.Size([80]) from checkpoint, the shape in current model is torch.Size([27]).

解决方法：去除预训练模型中该层的参数。

1.获取预训练模型中该层参数的名称。

报错信息中给出了参数名，也可读取模型参数

import torch
    model = torch.load(open(model_path, 'rb'), map_location='cpu')["state_dict"]
    for k,v in model.item():
        print(k)

2.删除预训练模型相关参数，再加载

model.pop('module.rpn.head.cls_logits.weight')
    model.pop('module.rpn.head.cls_logits.bias')
    new_model.load_state_dict(model, strict=False)

此处导入预训练模型时，一定要加上strict=False。参数strict缺省默认为True，一般深拷贝模型参数时，遇到unexpected\_keys或missing\_keys，会报错，设置strict=False则自动忽略这些keys，不做处理。此处pop参数之后，输出类别数不相同时，strict=True不会报错，但会拷贝预训练模型该权值（如下图）。设置strict=False后，则不会load相应参数权值。

![20210429140200893.PNG][]

附加：由于多GPU并行训练，模型参数会有"module"作前缀，建立空模型进行权值加载时，先去除预训练模型参数的"module"前缀，如下：

def delete(state_dict, prefix="module."):
        keys = sorted(state_dict.keys())
        if not all(key.startswith(prefix) for key in keys):
            return state_dict
        stripped_state_dict = OrderedDict()
        for key, value in state_dict.items():
            stripped_state_dict[key.replace(prefix, "")] = value
        return stripped_state_dict

[20210429140200893.PNG]: /images/20221021/e56439a2644f41c7bf011e0519662180.png