from typing import Tuple import torch from torch import nn import torchvision def create_effnetb3_model(num_classes: int = 101, 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_B3_Weights.DEFAULT transforms = weights.transforms() model = torchvision.models.efficientnet_b3(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=1536, out_features=num_classes) ) return model, transforms