pytorch：如何修改加载了预训练权重的模型的输入或输出--(原权重文件修改参数)

在使用pytorch的过程中，我们往往会使用官方发布的预训练模型，并在此基础上训练自己的模型。为了适配训练数据，有时候需要局部修改这类预训练模型的结构，本文将分别以修改输入的通道数和输出的分类数为例，讲解一种通用的方法来修练模型的结构。

加载模型

在修改之前需要加载预训练模型，这里以mobilenet v2为例

import torchvision.models as models
model = models.mobilenet_v2(pretrained=True)

修改模型结构

修改之前需要查看模型结构

print(model)

则可以看到一长串的模型输入，这里因为篇幅原因只截了开头部分和结尾部分

仔细观察输出的模型结构，卷积层（特别是括号中的features,classifier，(0) 等标志性词可以得知模型的第一层为:

model.features[0][0] = Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)

分类层为

model.classifier[1] = Linear(in_features=1280, out_features=1000, bias=True)

通过这些信息接下来就可以修改输入/输出了

修改模型输入

#输入为单通道

model.features[0][0] = Conv2d(1 ,32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)

#修改预训练模型权重的结构，使得模型可以使用修改后的预训练模型权重

#加载预训练模型
pre_trained_model = models.mobilenet_v2(pretrained=True)
#获取预训练权重文件的字典
pretrained_dict = pre_trained_model.state_dict()
#打印权重信息
print(pretrained_dict.items())
'''
打印显示为:
dict_items([('features.0.weight', tensor([[[[-2.8656e-03,  4.1653e-02,  5.7146e-02,  ...,  5.2015e-03,
           -5.7198e-03, -2.3688e-02],...
这里可以看到第一层的key为‘features.0.weight’,接下来就可以通过这个名称访问pretrained_dict中对应的权重
'''
#获取第一层权重
layer1 = pretrained_dict['features.0.0.weight']
#创建一个新的张量，这个张量后面将替代pretrain_dict中的第一层，以适应修改为单通道的模型
new = torch.zeros(32,1,3, 3)
#这里修改第一层
for i,output_channel in enumerate(layer1):
	# Grey = 0.299R + 0.587G + 0.114B, 这个公式参考了RGB图转灰度图的方式
    new[i] = 0.299 * output_channel[0] + 0.587 * output_channel[1] + 0.114 * output_channel[2]
#现在第一层的shape为（32，1，3，3）了
pretrained_dict['features.0.0.weight'] = new 
#修改模型结构
model.features[0][0] = nn.Conv2d(1, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
model.load_state_dict(pretrained_dict)

修改模型输出

这里修改输出的方式不像修改输入这么繁琐

fc_features = model.classifier[1].in_features
model.classifier[1] = nn.Linear(fc_features, 2)