Pytorch 預訓練班級數 Model

Question

我想使用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)

Answer 1

您必須更改相應 model 的最終線性層。

以re.net為例，當我們打印model時，我們看到最后一層是全連接層，如下圖：

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

因此，您必須將 model.fc 重新初始化為具有 512 個輸入特征和 2 個 output 特征的線性層：

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

對於其他型號，您可以在此處查看

要凍結 .network 的參數，您必須使用以下代碼：

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”是分類層的名稱。 其他模型並非如此。 您必須檢查 model 才能找到分類層的名稱。