oliverlevn/FoodMini

1import torch2import torchvision3 4from torch import nn5 6 7def create_effnetb2_model(num_classes:int=3, 8                          seed:int=42):9    """Creates an EfficientNetB2 feature extractor model and transforms.10 11    Args:12        num_classes (int, optional): number of classes in the classifier head. 13            Defaults to 3.14        seed (int, optional): random seed value. Defaults to 42.15 16    Returns:17        model (torch.nn.Module): EffNetB2 feature extractor model. 18        transforms (torchvision.transforms): EffNetB2 image transforms.19    """20    # Create EffNetB2 pretrained weights, transforms and model21    weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT22    transforms = weights.transforms()23    model = torchvision.models.efficientnet_b2(weights=weights)24 25    # Freeze all layers in base model26    for param in model.parameters():27        param.requires_grad = False28 29    # Change classifier head with random seed for reproducibility30    torch.manual_seed(seed)31    model.classifier = nn.Sequential(32        nn.Dropout(p=0.3, inplace=True),33        nn.Linear(in_features=1408, out_features=num_classes),34    )35    36    return model, transforms37