Spaces:
Sleeping
Sleeping
| """ | |
| Vision Transformer (ViT) transfer learning model. | |
| Uses timm library for pretrained ViT-B/16. | |
| Supports selective layer unfreezing and differential LR. | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| try: | |
| import timm | |
| TIMM_AVAILABLE = True | |
| except ImportError: | |
| TIMM_AVAILABLE = False | |
| print("Warning: timm not installed. ViT model requires: pip install timm") | |
| class CattleViT(nn.Module): | |
| """ | |
| Vision Transformer classifier using pretrained ViT from timm. | |
| Phase 1: Freeze all blocks, train only the classification head. | |
| Phase 2: Selectively unfreeze last N blocks for fine-tuning. | |
| """ | |
| def __init__( | |
| self, | |
| num_classes: int = 26, | |
| backbone: str = 'vit_base_patch16_224', | |
| pretrained: bool = True, | |
| freeze_backbone: bool = True, | |
| dropout: float = 0.1, | |
| ): | |
| super().__init__() | |
| if not TIMM_AVAILABLE: | |
| raise ImportError("timm is required for ViT. Install with: pip install timm") | |
| self.backbone_name = backbone | |
| # Load pretrained ViT | |
| self.model = timm.create_model( | |
| backbone, | |
| pretrained=pretrained, | |
| num_classes=0, # Remove classification head | |
| ) | |
| # Get feature dimension | |
| num_features = self.model.num_features | |
| # Custom classification head | |
| self.head = nn.Sequential( | |
| nn.LayerNorm(num_features), | |
| nn.Dropout(p=dropout), | |
| nn.Linear(num_features, 512), | |
| nn.GELU(), | |
| nn.Dropout(p=dropout * 0.5), | |
| nn.Linear(512, num_classes), | |
| ) | |
| if freeze_backbone: | |
| self.freeze_backbone() | |
| def freeze_backbone(self) -> None: | |
| """Freeze all backbone parameters.""" | |
| for param in self.model.parameters(): | |
| param.requires_grad = False | |
| def unfreeze_layers(self, layer_names: list[str] = None) -> None: | |
| """ | |
| Unfreeze specified layers for fine-tuning. | |
| Default: last 2 transformer blocks + norm layer. | |
| """ | |
| if layer_names is None: | |
| layer_names = ['blocks.10', 'blocks.11', 'norm'] | |
| 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 | |
| for param in self.head.parameters(): | |
| param.requires_grad = True | |
| def get_trainable_params(self) -> list: | |
| """Return list of all trainable parameters.""" | |
| params = [p for p in self.model.parameters() if p.requires_grad] | |
| params += [p for p in self.head.parameters() if p.requires_grad] | |
| return params | |
| 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 = [p for p in self.model.parameters() if p.requires_grad] | |
| head_params = list(self.head.parameters()) | |
| 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: | |
| features = self.model(x) | |
| return self.head(features) | |
| def from_config(config: dict) -> 'CattleViT': | |
| """Create model from config dict.""" | |
| arch = config.get('model', {}).get('architecture', {}) | |
| return CattleViT( | |
| num_classes=config.get('num_classes', 26), | |
| backbone=arch.get('backbone', 'vit_base_patch16_224'), | |
| pretrained=arch.get('pretrained', True), | |
| freeze_backbone=arch.get('freeze_backbone', True), | |
| ) | |