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

た入场券 2023-05-21 13:54 13阅读 0赞

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

## 加载模型 ##

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

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

## 修改模型结构 ##

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

print(model)

则可以看到一长串的模型输入，这里因为篇幅原因只截了开头部分和结尾部分  
![在这里插入图片描述][20200419104421323.png]  
仔细观察输出的模型结构，卷积层（特别是括号中的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)

[20200419104421323.png]: /images/20230521/eb6ef67cb4fd4a12bb099249a6866443.png