Pytorch 预训练模型中的类数

Posted 2023-02-16

【中文标题】Pytorch 预训练模型中的类数

【英文标题】：The Number of Classes in Pytorch Pretrained Model

【发布时间】：2021-10-29 00:08:20

【问题描述】：

我想使用 Pytorch 中的预训练模型在自己的数据集中进行图像分类，但是我应该如何在冻结特征提取层的参数的同时更改类的数量？

这些是我想要包括的模型：

resnet18 = models.resnet18(pretrained=True)
densenet161 = models.densenet161(pretrained=True)
inception_v3 = models.inception_v3(pretrained=True)
shufflenet_v2_x1_0 = models.shufflenet_v2_x1_0(pretrained=True)
mobilenet_v3_large = models.mobilenet_v3_large(pretrained=True)
mobilenet_v3_small = models.mobilenet_v3_small(pretrained=True)
mnasnet1_0 = models.mnasnet1_0(pretrained=True)
resnext50_32x4d = models.resnext50_32x4d(pretrained=True)
vgg16 = models.vgg16(pretrained=True)

提前非常感谢！

---

我添加的新代码：

import torch
from torchvision import models

class MyResModel(torch.nn.Module):
    def __init__(self):
        super(MyResModel, self).__init__()
        self.classifier = nn.Sequential(
            nn.Linear(512,256),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(256,3),
        )

    def forward(self, x):
        return self.classifier(x)

resnet18 = models.resnet18(pretrained=True)
resnet18.fc = MyResModel()

for param in resnet18.parameters():
    param.requires_grad_(False)

【参考方案1】：

您必须更改相应模型的最终线性层。

以resnet为例，当我们打印模型时，我们看到最后一层是全连接层，如下图：

(fc): Linear(in_features=512, out_features=1000, bias=True)

因此，您必须将 model.fc 重新初始化为具有 512 个输入特征和 2 个输出特征的线性层：

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

其他型号可以查看here

要冻结网络的参数，您必须使用以下代码：

for name, param in model.named_parameters():
    if 'fc' not in name:
        print(name, param.requires_grad)
        param.requires_grad=False

验证：

for name, param in model.named_parameters():  
    print(name,param.requires_grad)

请注意，对于此示例，“fc”是分类层的名称。这不是其他型号的情况。您必须检查模型才能找到分类层的名称。

【讨论】：

Thx，我现在试试，如果添加了多个fc层怎么办？我自己试了一下，不确定是否正确，我会立即更新问题，请您检查一下。

如果你的分类模块像 resnet18.fc = MyResModel() 一样在 'fc' 中，那么上面的代码应该可以工作。