小编Tim*_*Tim的帖子

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