Spaces:
Sleeping
Sleeping
| # ------------------------------------------------------------- | |
| # 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 |