model.py

import torch
import torchvision

from torch import nn


def model_efficientb3(out_feature:int=3,
                      p:int=0.3):
    """Creates an EfficientNetB2 feature extractor model and transforms.

    Args:
        num_classes (int, optional): number of classes in the classifier head. 
            Defaults to 3.
        seed (int, optional): random seed value. Defaults to 42.

    Returns:
        model (torch.nn.Module): EffNetB2 feature extractor model. 
        transforms (torchvision.transforms): EffNetB2 image transforms.
    """

    weights=torchvision.models.EfficientNet_B3_Weights.DEFAULT
    transform=weights.transforms()
    model=torchvision.models.efficientnet_b3(weights=weights)

    for params in model.parameters():
        params.requires_grad=False

    print(model.classifier)

    model.classifier=nn.Sequential(
        nn.Dropout(p=p,inplace=True),
        nn.Linear(in_features=1536,out_features=out_feature,bias=True)
    )

    print(f"the new classifier as per your request \n {model.classifier}")

    return model,transform