model.py

import torch
import torchvision
import torch.nn as nn


def create_effnetb2_model(num_classes: int = 3,
                          seed: int =42):

    # 1. Setup pretrained EffNetB2 weights
    effnetb2_weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT   # DEFAULT = best available weights.

    # 2. Get EffNetB2 transforms
    effnetb2_transforms = effnetb2_weights.transforms()

    # 3. Setup pretrained model instance
    effnetb2 = torchvision.models.efficientnet_b2(weights = effnetb2_weights)

    # 4. Freeze the base layers
    for param in effnetb2.parameters():
        param.requires_grad = False
        
    # 5. Change classifier head with random seed for reproducibility.    
    torch.manual_seed(seed)
    effnetb2.classifier = nn.Sequential(
        nn.Dropout(p = 0.3, inplace = True),
        nn.Linear(in_features = 1408, out_features = num_classes, bias = True)
        )
    
    return effnetb2, effnetb2_transforms