pytorch 模型部分参数的加载

方法一：保留同名模块进行加载

如果对预训练模型的结构进行了一些改动，在训练的开始前希望加载未改动部分的参数，如将resnet18的第一层卷积层conv1的输入由3通道改为6通道的new_conv1，将分类层fc的1000类输出改为2类输出的new_fc，注意：要改一下名字与原来的不同。

导入模型

myNet=resnet18()
   
     
     
     
     

然后就加载模型的参数，参考pytorch 如何加载部分预训练模型

pretrained_dict=torch.load(model_weight)
model_dict=myNet.state_dict()
# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
myNet.load_state_dict(model_dict)

也可以通过pretrained model.state_dict()提取需要的模型参数。

方法二：使用torch.load_dict加载（使用strict参数）

myNet.load_state_dict(torch.load(model_weight),
                      strict=False) 

这一句话就搞定了，key相同（key可以理解为模块名字）的加载进去，不相同的就丢弃掉了。 注意，若你更改了比如某个conv的output_channel, 此时key还是相同的，当你使用load_state_dict加载时就会报错。以下是使用coco（184类）训练的deeplab, 尝试加载预训练权重到用于训练voc数据集时（只有21类）就会出现：

Error(s) in loading state_dict for DataParallelWithCallback:
1. size mismatch for module.decoder.output.7.weight: copying a param with shape torch.Size([184, 256, 1, 1])from checkpoint,
the shape in current model is torch.Size([91, 256, 1, 1]).
2. size mismatch for module.decoder.output.7.bias: copying a param with shape torch.Size([184]) from checkpoint, the shape in current model is torch.Size([91]).