Spaces:
Sleeping
Sleeping
| """ | |
| ResNet transfer learning model. | |
| Wraps torchvision ResNet50 with custom classifier head. | |
| Supports two-phase training: frozen backbone → partial fine-tuning. | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import torchvision.models as models | |
| from torchvision.models import ResNet50_Weights, ResNet18_Weights | |
| class CattleResNet(nn.Module): | |
| """ | |
| ResNet-based classifier with transfer learning support. | |
| Phase 1: Freeze all backbone layers, train only the classifier head. | |
| Phase 2: Unfreeze specified layers for fine-tuning. | |
| """ | |
| def __init__( | |
| self, | |
| num_classes: int = 26, | |
| backbone: str = 'resnet50', | |
| pretrained: bool = True, | |
| freeze_backbone: bool = True, | |
| dropout: float = 0.3, | |
| ): | |
| super().__init__() | |
| self.backbone_name = backbone | |
| # Load pretrained backbone | |
| if backbone == 'resnet50': | |
| weights = ResNet50_Weights.IMAGENET1K_V1 if pretrained else None | |
| self.model = models.resnet50(weights=weights) | |
| elif backbone == 'resnet18': | |
| weights = ResNet18_Weights.IMAGENET1K_V1 if pretrained else None | |
| self.model = models.resnet18(weights=weights) | |
| else: | |
| raise ValueError(f"Unsupported backbone: {backbone}") | |
| # Replace classifier head | |
| num_features = self.model.fc.in_features | |
| self.model.fc = nn.Sequential( | |
| nn.Dropout(p=dropout), | |
| nn.Linear(num_features, 512), | |
| nn.ReLU(inplace=True), | |
| nn.Dropout(p=dropout * 0.5), | |
| nn.Linear(512, num_classes), | |
| ) | |
| # Freeze backbone if requested | |
| if freeze_backbone: | |
| self.freeze_backbone() | |
| def freeze_backbone(self) -> None: | |
| """Freeze all layers except the classifier head.""" | |
| for name, param in self.model.named_parameters(): | |
| if 'fc' not in name: | |
| param.requires_grad = False | |
| def unfreeze_layers(self, layer_names: list[str] = None) -> None: | |
| """ | |
| Unfreeze specified layers for fine-tuning. | |
| Default: unfreeze layer3, layer4, and fc. | |
| """ | |
| if layer_names is None: | |
| layer_names = ['layer3', 'layer4', 'fc'] | |
| for name, param in self.model.named_parameters(): | |
| if any(ln in name for ln in layer_names): | |
| param.requires_grad = True | |
| def unfreeze_all(self) -> None: | |
| """Unfreeze all parameters.""" | |
| for param in self.model.parameters(): | |
| param.requires_grad = True | |
| def get_trainable_params(self) -> list: | |
| """Return list of trainable parameters.""" | |
| return [p for p in self.model.parameters() if p.requires_grad] | |
| def get_param_groups(self, base_lr: float, fine_tune_lr: float) -> list[dict]: | |
| """ | |
| Return parameter groups with different learning rates. | |
| Backbone gets fine_tune_lr, head gets base_lr. | |
| """ | |
| backbone_params = [] | |
| head_params = [] | |
| for name, param in self.model.named_parameters(): | |
| if not param.requires_grad: | |
| continue | |
| if 'fc' in name: | |
| head_params.append(param) | |
| else: | |
| backbone_params.append(param) | |
| param_groups = [] | |
| if backbone_params: | |
| param_groups.append({'params': backbone_params, 'lr': fine_tune_lr}) | |
| if head_params: | |
| param_groups.append({'params': head_params, 'lr': base_lr}) | |
| return param_groups | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return self.model(x) | |
| def from_config(config: dict) -> 'CattleResNet': | |
| """Create model from config dict.""" | |
| arch = config.get('model', {}).get('architecture', {}) | |
| return CattleResNet( | |
| num_classes=config.get('num_classes', 26), | |
| backbone=arch.get('backbone', 'resnet50'), | |
| pretrained=arch.get('pretrained', True), | |
| freeze_backbone=arch.get('freeze_backbone', True), | |
| ) | |