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

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

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相应参数权值。

附加：由于多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