Pytorch模型微调fine-tune详解

def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad = False

import torchvision.models as models
# 冻结参数的梯度
feature_extract = True
model = models.resnet18(pretrained=True)
set_parameter_requires_grad(model, feature_extract)
# 修改模型, 输出通道4， 此时，fc层就被随机初始化了，但是其他层依然保存着预训练得到的参数。
model.fc = nn.Linear(in_features=512, out_features=4, bias=True)

# 8 更改池化层
class AdaptiveConcatPool2d(nn.Module):
    def __init__(self, size=None):
        super().__init__()
        size = size or (1, 1) # 池化层的卷积核大小，默认值为（1，1）
        self.pool_one = nn.AdaptiveAvgPool2d(size) # 池化层1
        self.pool_two = nn.AdaptiveAvgPool2d(size) # 池化层2
 
    def forward(self, x):
        return torch.cat([self.pool_one(x), self.pool_two(x), 1]) # 连接两个池化层
 
 
# 7 迁移学习：拿到一个成熟的模型，进行模型微调
def get_model():
    model_pre = models.resnet50(pretrained=True) # 获取预训练模型
    # 冻结预训练模型中所有的参数
    for param in model_pre.parameters():
        param.requires_grad = False
    # 微调模型：替换ResNet最后的两层网络，返回一个新的模型
    model_pre.avgpool = AdaptiveConcatPool2d() # 池化层替换
    model_pre.fc = nn.Sequential(
        nn.Flatten(), # 所有维度拉平
        nn.BatchNorm1d(4096), # 256 x 6 x 6   ——> 4096
        nn.Dropout(0.5),  # 丢掉一些神经元
        nn.Linear(4096, 512),  # 线性层的处理
        nn.ReLU(), # 激活层
        nn.BatchNorm1d(512), # 正则化处理
        nn.Linear(512,2),
        nn.LogSoftmax(dim=1), # 损失函数
    )
 
    return

if freeze_layers:
       for name, param in model.named_parameters():
           # 除最后的全连接层外，其他权重全部冻结
           if "fc" not in name:
               param.requires_grad_(False)
 
pg = [p for p in model.parameters() if p.requires_grad]
optimizer = optim.SGD(pg, lr=0.01, momentum=0.9, weight_decay=4E-5)

if Freeze_Train:
   for param in model.backbone.parameters():
       param.requires_grad = False

count = 0
for layer in model.children():
   count = count + 1
   if count < 10:
       for param in layer.parameters():
           param.requires_grad = False
 
# 然后将需要训练的参数传入优化器，也就是过滤掉被冻结的参数。
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=LR)

model_dict = pretrained_model.state_dict()
pretrained_dict={k: v for k, v in pretrained_dict.items() if k[7:] in model_dict}
model_dict.update(pretrained_dict)

import torch.nn as nn
import torch
class AlexNet(nn.Module):
   def __init__(self):
       super(AlexNet, self).__init__()
       self.features=nn.Sequential(
           nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),  # 使用卷积层，输入为3，输出为64，核大小为11，步长为4
           nn.ReLU(inplace=True),  # 使用激活函数
           nn.MaxPool2d(kernel_size=3, stride=2),  # 使用最大池化，这里的大小为3，步长为2
           nn.Conv2d(64, 192, kernel_size=5, padding=2), # 使用卷积层，输入为64，输出为192，核大小为5，步长为2
           nn.ReLU(inplace=True),# 使用激活函数
           nn.MaxPool2d(kernel_size=3, stride=2), # 使用最大池化，这里的大小为3，步长为2
           nn.Conv2d(192, 384, kernel_size=3, padding=1), # 使用卷积层，输入为192，输出为384，核大小为3，步长为1
           nn.ReLU(inplace=True),# 使用激活函数
           nn.Conv2d(384, 256, kernel_size=3, padding=1),# 使用卷积层，输入为384，输出为256，核大小为3，步长为1
           nn.ReLU(inplace=True),# 使用激活函数
           nn.Conv2d(256, 256, kernel_size=3, padding=1),# 使用卷积层，输入为256，输出为256，核大小为3，步长为1
           nn.ReLU(inplace=True),# 使用激活函数
           nn.MaxPool2d(kernel_size=3, stride=2),  # 使用最大池化，这里的大小为3，步长为2
      )
       self.avgpool=nn.AdaptiveAvgPool2d((6, 6))
       self.classifier=nn.Sequential(
           nn.Dropout(),# 使用Dropout来减缓过拟合
           nn.Linear(256 * 6 * 6, 4096),   # 全连接，输出为4096
           nn.ReLU(inplace=True),# 使用激活函数
           nn.Dropout(),# 使用Dropout来减缓过拟合
           nn.Linear(4096, 4096),  # 维度不变，因为后面引入了激活函数，从而引入非线性
           nn.ReLU(inplace=True),  # 使用激活函数
           nn.Linear(4096, 1000),   #ImageNet默认为1000个类别，所以这里进行1000个类别分类
      )
     
 
   def forward(self, x):
       x=self.features(x)
       x=self.avgpool(x)
       x=torch.flatten(x, 1)
       x=self.classifier(x)
       return x
 
def alexnet(num_classes, device, pretrained_weights=""):
   net=AlexNet()  # 定义AlexNet
   if pretrained_weights:  # 判断预训练模型路径是否为空，如果不为空则加载
       net.load_state_dict(torch.load(pretrained_weights,map_location=device))
   
   num_fc=net.classifier[6].in_features  # 获取输入到全连接层的输入维度信息
   net.classifier[6]=torch.nn.Linear(in_features=num_fc, out_features=num_classes) # 根据数据集的类别数来指定最后输出的out_features数目
   return net