src/model.py

import torch
import torch.nn as nn
import torchvision
from src.logger import global_logger as logger
from torchvision.models import resnet50, ResNet50_Weights

def resnet_model(num_classes: int = 4, seed: int = 42):
    # Load pretrained ResNet18 model
    weights = ResNet50_Weights.DEFAULT
    model = resnet50(weights=weights)

    # Freeze the parameters of the pretrained model
    for param in model.parameters():
        param.requires_grad = False

    #logger.info("Model initialized with frozen ResNet18 backbone and new fully connected layers.")

    # Replace the final fully connected layer with a new one
    torch.manual_seed(seed)
    model.fc = nn.Sequential(
        nn.Dropout(p=0.3, inplace=True),
        nn.Linear(in_features=model.fc.in_features, out_features=num_classes),
    )

    # Define the transforms using the predefined transforms from weights
    transforms = weights.transforms()

    return model, transforms

# Example usage
model, transforms = resnet_model(num_classes=4)