model.py

# -*- coding: utf-8 -*-
"""
Created on Thu Feb  8 13:48:19 2024

@author: firis
"""

import torch
import torchvision
from torch import nn

def create_eff_model(num_classes:int=3,seed:int=42):
    
    weights=torchvision.models.EfficientNet_B1_Weights.DEFAULT
    transforms=weights.transforms()
    model =torchvision.models.efficientnet_b1(weights=weights)

    # Freeze all layers in base model
    for param in model.parameters():
        param.requires_grad = False
        
    torch.manual_seed(seed)
    
    model.classifier = nn.Sequential(
      nn.Dropout(p=0.3, inplace=True),
      nn.Linear(in_features=1280, out_features=num_classes))
    
    return model, transforms