glisicstefan/age-estimation-resnet50

1from torchvision.models import ResNet50_Weights, resnet502from torch import nn3 4def create_resnet50(out_features: int = 1, freeze: bool = True):5    # 1. Uzmi predtrenirane težine i transformacije6    weights = ResNet50_Weights.DEFAULT7    transforms = weights.transforms()8 9    # 2. Učitaj model10    model = resnet50(weights=weights)11 12    # 3. Zamrzni ili otključaj bazu (backbone)13    # Ako je freeze=True, gradijenti se ne računaju (zamrznuto)14    for param in model.parameters():15        param.requires_grad = not freeze16 17    # 4. Izmeni FC sloj (on je uvek otključan po defaultu jer je nov)18    # in_features je 2048 za ResNet5019    model.fc = nn.Linear(in_features=2048, out_features=out_features)20 21    return model, transforms