How to add additional layer( inside the features module) in pre-trained EfficientNetB0 model?

[sashika-himali]

Could you please guide me on how to add a multi-head attention layer after model. features[5:] . I want to add this layer inside the features .
import torch
import torchvision
weights = torchvision. models.EfficientNet_B0_Weights.DEFAULT
model = torchvision.models.efficientnet_b0(weights=weights).to(device)

[cthugha10]

Try just create another nn.mudule which includes the pre-trained model and your layer?

[adhilcodes]

class EfficientNet_bo(nn.Module):
    def __init__(self, num_classes = 2):
        super().__init__()
        self.model = EfficientNet.from_pretrained('efficientnet-b0')

        for param in self.model.parameters():
            param.requires_grad = False

        # Extract number of features from last layer
        in_features = self.model._fc.in_features

        self.model._fc = nn.Sequential(
            nn.Linear(in_features, 512),
            nn.ReLU(inplace=True),
            nn.Linear(512, 256),
            nn.ReLU(inplace=True),
            nn.Linear(256, num_classes),
            nn.Sigmoid() 
        )

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




    Here is a sample. Customise in the way you want it