How to add the last classification layer in EfficieNet pre-trained model in Pytorch?

Question

I'm using the EfficientNet pre-trained model for my image classification project in Pytorch, and my purpose is to change the number of classes which is initially 1000 to 4. However, for that when I try adding a model._fc layer, I keep on seeing this error "EfficientNet' object has no attribute 'classifier". Here is my code (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)

The situation is different when I add model._fc to the end of the Resnet part, it clearly changes the number of output classes to 4 in Resnet-18. Here is the code for that:

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, bias=True)

The solution is available for TensorFlow and Keras, and I would really appreciate it if anyone could help me with that in PyTorch.

Regards,

Far

Answer 1

EfficentNet class doesn't have attribute classifier, you need to change in_features=model.classifier.in_features to in_features=model._fc.in_features.

import torchvision.models as models

NUM_CLASSES = 4

#EfficientNet
from efficientnet_pytorch import EfficientNet
efficientnet = EfficientNet.from_pretrained('efficientnet-b3')
efficientnet ._fc= torch.nn.Linear(in_features=efficientnet._fc.in_features, out_features=NUM_CLASSES, bias=True)

#mobilenet_v2
mobilenet = models.mobilenet_v2(pretrained=True)
mobilenet.classifier = nn.Sequential(nn.Dropout(p=0.2, inplace=False),
                  nn.Linear(in_features=mobilenet.classifier[1].in_features, out_features=NUM_CLASSES, bias=True))

#inception_v3
inception = models.inception_v3(pretrained=True)
inception.fc =  nn.Linear(in_features=inception.fc.in_features, out_features=NUM_CLASSES, bias=True)