Spaces:
Sleeping
Sleeping
File size: 2,219 Bytes
58886b6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | # -------------------------------------------------------------
# This script builds and returns your model architecture.
# You’re using ResNet-18 (a pre-trained CNN) as the backbone.
#
# Key idea:
# - Load a pre-trained ResNet-18 (trained on ImageNet)
# - Optionally freeze its layers (so they don’t update)
# - Replace its final layer with your custom classifier
# - Move the model to GPU or CPU
# -------------------------------------------------------------
import torch
import torch.nn as nn
from torchvision import models
def build_model(num_classes: int, freeze_backbone: bool, device: torch.device) -> nn.Module:
# ---------------------------------------------------------
# 1️⃣ Load a pre-trained ResNet-18
# ---------------------------------------------------------
try:
weights = models.ResNet18_Weights.IMAGENET1K_V1 # For newer torchvision versions
model = models.resnet18(weights=weights)
# Older torchvision fallback
except Exception:
model = models.resnet18(pretrained=True)
# ---------------------------------------------------------
# 2️⃣ Optionally freeze the backbone (for transfer learning)
# ---------------------------------------------------------
# If freeze_backbone=True, only the last (new) layer will train
if freeze_backbone:
for p in model.parameters(): # No gradients, no updates for any layer you froze
p.requires_grad = False
# ---------------------------------------------------------
# 3️⃣ Replace the final fully-connected layer (the classifier)
# ---------------------------------------------------------
# The default ResNet18 outputs 1000 classes (ImageNet)
# You’re replacing it with your own number of output classes
in_feats = model.fc.in_features # Get input features of last layer
model.fc = nn.Linear(in_feats, num_classes) # Replace with new layer
# ---------------------------------------------------------
# 4️⃣ Move model to selected device (GPU or CPU)
# ---------------------------------------------------------
model.to(device)
return model |