deep-learning - 如何在 Pytorch 中添加 EfficieNet 预训练模型的最后一个分类层？

Question

我正在使用EfficientNet我在 Pytorch 中为我的图像分类项目预训练了模型，我的目的是将最初的 1000 个类的数量更改为 4 个。 然而，为此，当我尝试添加 model._fc 层时，我不断看到此错误“EfficientNet”对象没有属性“分类器”。这是我的代码(Config.NUM_CLASSES = 4):

elif Config.MODEL_NAME == 'efficientnet-b3':
      
    from efficientnet_pytorch import EfficientNet
    model = EfficientNet.from_pretrained('efficientnet-b3')
    model._fc= torch.nn.Linear(in_features=model.classifier.in_features, **out_features=Config.NUM_CLASSES**, bias=True)

当我将 model._fc 添加到 Resnet 部分的末尾时，情况有所不同，它明显将 Resnet-18 中的输出类数量更改为 4。这是代码:

if Config.MODEL_NAME == 'resnet18': model = models.resnet50(pretrained=True) model.fc = torch.nn.Linear(in_features=model.fc.in_features, out_features=Config.NUM_CLASSES,偏差=真)

solution适用于 TensorFlow 和 Keras，如果有人能在 PyTorch 中帮助我，我将非常感激。

问候，

远

Answer 1

Torchvision >= 0.11 包括 EfficientNet，并且它确实具有分类器属性。获取 in_features :

import torchvision

model = torchvision.models.efficientnet_b5()
num_ftrs = model.classifier[1].in_features