model.py

import torch
import torchvision
from torch import nn

def create_effnetb2_model(num_classes : int = 3, # default output classes = 3
                          seed : int = 42):
  # 1,2,3 create EffNetB2 pretrained weights, transforms, and model
  weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
  transforms = weights.transforms()
  model = torchvision.models.efficientnet_b2(weights = weights)
  
  # 4. Freeze all base layers in the model
  for param in model.parameters():
    param.requires_grad = False

  # 5. Change classifier head with random seet for reproducibility
  torch.manual_seed(seed)
  model.classifier = nn.Sequential(
      nn.Dropout(p=0.3, inplace =True),
      nn.Linear(in_features = 1408, out_features = num_classes, bias = True))

  return model, transforms