import torch import torchvision from torch import nn def create_effnetb4_model(num_classes: int = 101, seed: int = 42): """ Creates an EfficientNet-B4 feature extractor model and its preprocessing transforms. Args: num_classes (int, optional): Number of output classes. Defaults to 101 (Food-101 dataset). seed (int, optional): Random seed for reproducibility. Defaults to 42. Returns: model (torch.nn.Module): EfficientNet-B4 feature extractor model. transforms (torchvision.transforms.Compose): Corresponding image transforms. """ # Use pretrained EfficientNet-B4 weights weights = torchvision.models.EfficientNet_B4_Weights.DEFAULT transforms = weights.transforms() # Initialize model model = torchvision.models.efficientnet_b4(weights=weights) # Freeze feature extractor layers for param in model.parameters(): param.requires_grad = False # Set seed for reproducibility torch.manual_seed(seed) # Replace classifier head model.classifier = nn.Sequential( nn.Dropout(p=0.3, inplace=True), nn.Linear(in_features=1792, out_features=num_classes) ) return model, transforms