model.py

from typing import Tuple

import torch
from torch import nn
import torchvision

def create_effnetb2_model(num_classes: int = 3,
                          seed: int = 4,
                         ) -> Tuple[nn.Module, torchvision.transforms.Compose]:
    """Create an EfficientNetB2 feature extractor model and transforms.
    
    Args:
      num_classes: Number of classes to use for classification (default 3).
      seed: Random seed for reproducibility (default 4).
      
    Returns:
      A tuple (model, transforms) of the model and its image transforms.
    """
    weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
    transforms = weights.transforms()
    model = torchvision.models.efficientnet_b2(weights=weights)
    
    # Freeze parameters below the head
    for param in model.parameters():
        param.requires_grad = False
    # Replace the classifier head with one of appropriate size for the problem
    torch.manual_seed(seed)
    model.classifier = nn.Sequential(
        nn.Dropout(p=0.3, inplace=True),
        nn.Linear(in_features=1408, out_features=len(class_names))
    )
    return model, transforms