model.py

import torch
from torch import nn
import torchvision
from torchvision import transforms,models
def create_effnetb2_model(num_classes:int):
    """ 
    Function that creates the EffnetB2 model, freeze the features and update the classifier
    with the correspondinf number of classes.

    Return:
    1. The EffnetB2 model
    2. Transforms for the model
    3. Model summary
    """
    # Get weights
    weights_EffnetB2 = models.EfficientNet_B2_Weights.DEFAULT
    # Get trasnforms
    transform_b2 = weights_EffnetB2.transforms()
    # Create the model
    effnet_b2 = models.efficientnet_b2(weights=weights_EffnetB2)

    # Update the architecture:
    for param in effnet_b2.features.parameters():
        param.requires_grad = False
    
    effnet_b2.classifier = nn.Sequential(
        nn.Dropout(p=0.3,inplace=True),
        nn.Linear(in_features=1408,out_features=num_classes))
        
    return effnet_b2, transform_b2