model.py

import torch
import torchvision
from torch import nn

def create_effnetb2_model(num_classes: int=3, seed: int=42, device:torch.device="cuda" if torch.cuda.is_available() else "cpu"):

  """Creates an EfficientNetB2 feature extractor model and transforms.

  Args:
      num_classes (int, optional): number of classes in the classifier head.
          Defaults to 3.
      seed (int, optional): random seed value. Defaults to 42.

  Returns:
      model (torch.nn.Module): EffNetB2 feature extractor model.
      transforms (torchvision.transforms): EffNetB2 image transforms.
  """

  #set_seeds(seed)

  weights=torchvision.models.EfficientNet_B2_Weights.DEFAULT
  effnetb2_transforms=weights.transforms()
  effnetb2_model=torchvision.models.efficientnet_b2(weights=weights).to(device)

  for param in effnetb2_model.parameters():
    param.requires_grad=False

  effnetb2_model.classifier = nn.Sequential(
      nn.Dropout(p=0.3, inplace=True),
      nn.Linear(in_features=1408, out_features=num_classes)
  ).to(device)
  return effnetb2_model, effnetb2_transforms