model.py

import torchvision
from torch import nn

def create_effnet_b1(num_classes:int=5, seed:int=42):
    # 1. Get the base mdoel with pretrained weights and send to target device
    weights = torchvision.models.EfficientNet_B1_Weights.DEFAULT
    transforms = weights.transforms()
    model = torchvision.models.efficientnet_b1(weights=weights)#.to(device)

    # 2. Freeze the base model layers
    for param in model.features.parameters():
        param.requires_grad = False

    # 3. Change the classifier head
    model.classifier = nn.Sequential(
        nn.Dropout(p=0.2),
        nn.Linear(in_features=1280, out_features=5)
    )#.to(device)

    # 5. Give the model a name
    model.name = "effnet_b1"
    print(f"[INFO] Created new {model.name} model.")
    return model, transforms