如何修改预训练模型的部分结构，冻结部分层

以ResNet50做分类为例，加载预训练模型

import torch
import torchvision.models as models

#pretrained=True加载预训练模型参数，会下载到指定位置，为False则不加载参数
model = models.resnet50(pretrained=True)

 可以在终端直接运行model,查看模型结构

model

运行结果为：

 

 可以看到模型最后一层为fc层，in_features=2048, out_features=1000

1.针对自己的数据集的类别num_classes往往需要修改最后一层fc

model.fc = nn.Linear(2048, num_classes)

#如果不知道infeature = 2048, 可以通过如下方式获取：
fc_in_features = model.fc.in_features
model.fc = nn.Linear(fc_in_features , num_classes)

2.如果想删除某些层，可以通过恒等映射使输入等于输出，就相当于删除这层了

如修改conv1，并删除bn1、relu层

#修改conv1为3个3*3的卷积层
model.conv1 = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
        )

#恒等映射
model.bn1 = nn.Identity()
model.relu = nn.Identity()

 修改后再次查看模型结构，可以看到结构已经被修改

 3.冻结某些层，只需将这些层的参数的 requires_grad属性设置为 False即可。

如除了fc层，其他层都冻结：

# 冻结除全连接层以外的所有层
for name, param in model.named_parameters():
    if name != 'fc.weight' and name != 'fc.bias':
        param.requires_grad = False

可以通过以下代码查看是否被冻结

# 打印每个参数及其 requires_grad 属性
for name, param in model.named_parameters():
    print(f'{name}: {param.requires_grad}')

输出结果：