import torch import torch.nn as nn import timm from typing import Dict, Optional, Union, List import logging # Import configurations from separate file from .model_configs import MODEL_REGISTRY, TRAINING_CONFIGS, get_model_family, FINETUNE_CONFIGS from .base_model import ModelWithFeatures logger = logging.getLogger(__name__) class MultiTaskHead(nn.Module): """Multi-task head for decade, cluster, and device type prediction""" def __init__( self, in_features: int, num_decade_classes: int = 5, num_cluster_classes: int = 10, num_device_classes: int = 2, # phone or calculator hidden_dim: int = 512, dropout_rate: float = 0.3 ): super().__init__() self.num_decade_classes = num_decade_classes self.num_cluster_classes = num_cluster_classes self.num_device_classes = num_device_classes # Shared feature extractor self.shared_features = nn.Sequential( nn.Linear(in_features, hidden_dim), nn.ReLU(), nn.Dropout(dropout_rate), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Dropout(dropout_rate) ) # Task-specific heads self.decade_head = nn.Sequential( nn.Linear(hidden_dim, hidden_dim // 2), nn.ReLU(), nn.Dropout(dropout_rate), nn.Linear(hidden_dim // 2, num_decade_classes) ) self.cluster_head = nn.Sequential( nn.Linear(hidden_dim, hidden_dim // 2), nn.ReLU(), nn.Dropout(dropout_rate), nn.Linear(hidden_dim // 2, num_cluster_classes) ) # New device type head (phone vs calculator) self.device_head = nn.Sequential( nn.Linear(hidden_dim, hidden_dim // 2), nn.ReLU(), nn.Dropout(dropout_rate), nn.Linear(hidden_dim // 2, num_device_classes) ) # Initialize weights self._initialize_weights() def _initialize_weights(self): """Initialize weights properly""" for m in [self.shared_features, self.decade_head, self.cluster_head, self.device_head]: for layer in m: if isinstance(layer, nn.Linear): nn.init.kaiming_normal_(layer.weight) nn.init.constant_(layer.bias, 0) def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]: """Forward pass through multi-task head""" shared_features = self.shared_features(x) decade_logits = self.decade_head(shared_features) cluster_logits = self.cluster_head(shared_features) device_logits = self.device_head(shared_features) return { 'decade': decade_logits, 'cluster': cluster_logits, 'device': device_logits } class MultiTaskModel(nn.Module): """Wrapper to convert single-task model to multi-task""" def __init__( self, backbone: nn.Module, num_decade_classes: int = 5, num_cluster_classes: int = 10, num_device_classes: int = 2, hidden_dim: int = 512, dropout_rate: float = 0.3 ): super().__init__() self.backbone = backbone # Get the number of features from the backbone # This works for most timm models if hasattr(backbone, 'num_features'): in_features = backbone.num_features elif hasattr(backbone, 'classifier'): if isinstance(backbone.classifier, nn.Linear): in_features = backbone.classifier.in_features else: # For more complex classifiers, take the last linear layer in_features = None for module in reversed(list(backbone.classifier.modules())): if isinstance(module, nn.Linear): in_features = module.in_features break if in_features is None: raise ValueError("Could not determine backbone output features") else: raise ValueError("Could not determine backbone output features") # Replace the classifier with identity to get features if hasattr(backbone, 'classifier'): backbone.classifier = nn.Identity() elif hasattr(backbone, 'fc'): backbone.fc = nn.Identity() else: # Try to find and replace the last linear layer for name, module in backbone.named_modules(): if isinstance(module, nn.Linear) and 'classifier' in name.lower(): setattr(backbone, name.split('.')[-1], nn.Identity()) break # Create multi-task head self.multitask_head = MultiTaskHead( in_features=in_features, num_decade_classes=num_decade_classes, num_cluster_classes=num_cluster_classes, num_device_classes=num_device_classes, hidden_dim=hidden_dim, dropout_rate=dropout_rate ) logger.info(f"Created multi-task model with {in_features} backbone features") def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]: """Forward pass""" features = self.backbone(x) return self.multitask_head(features) def predict(self, x: torch.Tensor) -> Dict[str, torch.Tensor]: """Make predictions with softmax""" with torch.no_grad(): logits = self.forward(x) return { 'decade': torch.softmax(logits['decade'], dim=1), 'cluster': torch.softmax(logits['cluster'], dim=1), 'device': torch.softmax(logits['device'], dim=1) } class MultiTaskLoss(nn.Module): """Combined loss for multi-task learning""" def __init__( self, decade_weight: float = 1.0, cluster_weight: float = 1.0, device_weight: float = 1.0, loss_type: str = 'cross_entropy', loss_params: Optional[Dict] = None ): super().__init__() self.decade_weight = decade_weight self.cluster_weight = cluster_weight self.device_weight = device_weight # Import losses module to access all loss functions from ..training.losses import get_loss_function # Create loss functions for each task loss_params = loss_params or {} self.decade_criterion = get_loss_function(loss_type, **loss_params) self.cluster_criterion = get_loss_function(loss_type, **loss_params) self.device_criterion = get_loss_function(loss_type, **loss_params) def forward( self, predictions: Dict[str, torch.Tensor], targets: Dict[str, torch.Tensor] ) -> Dict[str, torch.Tensor]: """Calculate combined loss""" decade_loss = self.decade_criterion(predictions['decade'], targets['decade']) cluster_loss = self.cluster_criterion(predictions['cluster'], targets['cluster']) device_loss = self.device_criterion(predictions['device'], targets['device']) total_loss = (self.decade_weight * decade_loss + self.cluster_weight * cluster_loss + self.device_weight * device_loss) return { 'total_loss': total_loss, 'decade_loss': decade_loss, 'cluster_loss': cluster_loss, 'device_loss': device_loss } class ModelFactory: """Factory class for creating different model architectures""" @classmethod def create_model( cls, model_name: str, num_classes: Union[int, Dict[str, int]] = 5, # Can be int or dict for multi-task pretrained: bool = True, checkpoint_path: Optional[str] = None, return_features: bool = False, multi_task: bool = False, multitask_config: Optional[Dict] = None ) -> nn.Module: """ Create a model instance Args: model_name: Name of the model architecture num_classes: Number of output classes (int) or dict with task names for multi-task pretrained: Whether to use pretrained weights checkpoint_path: Path to load checkpoint from return_features: Wrap model to return features multi_task: Whether to create multi-task model multitask_config: Configuration for multi-task head Returns: Model instance """ if model_name not in MODEL_REGISTRY: raise ValueError(f"Unknown model: {model_name}. Available models: {list(MODEL_REGISTRY.keys())}") timm_model_name = MODEL_REGISTRY[model_name] if multi_task: # Multi-task model creation if isinstance(num_classes, dict): num_decade_classes = num_classes.get('decade', 5) num_cluster_classes = num_classes.get('cluster', 10) num_device_classes = num_classes.get('device', 2) else: # Assume single number is for decades, cluster and device classes need to be specified num_decade_classes = num_classes num_cluster_classes = multitask_config.get('num_cluster_classes', 10) if multitask_config else 10 num_device_classes = multitask_config.get('num_device_classes', 2) if multitask_config else 2 # Create backbone with dummy classifier (will be replaced) backbone = timm.create_model( timm_model_name, pretrained=pretrained, num_classes=1000 # Use original pretrained classes initially ) # Create multi-task wrapper multitask_config = multitask_config or {} model = MultiTaskModel( backbone=backbone, num_decade_classes=num_decade_classes, num_cluster_classes=num_cluster_classes, num_device_classes=num_device_classes, hidden_dim=multitask_config.get('hidden_dim', 512), dropout_rate=multitask_config.get('dropout_rate', 0.3) ) logger.info(f"Created multi-task model: {model_name}") logger.info(f"Decade classes: {num_decade_classes}, Cluster classes: {num_cluster_classes}, Device classes: {num_device_classes}") else: # Single-task model creation (original behavior) if isinstance(num_classes, dict): num_classes = num_classes.get('decade', 5) # Default to decade task model = timm.create_model( timm_model_name, pretrained=pretrained, num_classes=num_classes ) # Wrap with feature extractor if requested if return_features: model = ModelWithFeatures(model, num_classes=num_classes) logger.info(f"Created single-task model: {model_name}") logger.info(f"Number of parameters: {sum(p.numel() for p in model.parameters()):,}") # Load checkpoint if provided if checkpoint_path: try: checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=True) except: checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False) if 'model_state_dict' in checkpoint: try: model.load_state_dict(checkpoint['model_state_dict']) logger.info(f"Loaded checkpoint from {checkpoint_path}") except RuntimeError as e: logger.warning(f"Could not load full checkpoint due to architecture mismatch: {e}") logger.info("Attempting to load compatible layers only...") # Load compatible layers only model_dict = model.state_dict() checkpoint_dict = checkpoint['model_state_dict'] # Filter compatible layers compatible_dict = { k: v for k, v in checkpoint_dict.items() if k in model_dict and model_dict[k].shape == v.shape } model_dict.update(compatible_dict) model.load_state_dict(model_dict) logger.info(f"Loaded {len(compatible_dict)} compatible layers from checkpoint") else: model.load_state_dict(checkpoint) return model @classmethod def create_multitask_loss( cls, decade_weight: float = 1.0, cluster_weight: float = 1.0, device_weight: float = 1.0, loss_type: str = 'cross_entropy', loss_params: Optional[Dict] = None ) -> MultiTaskLoss: """Create multi-task loss function""" return MultiTaskLoss( decade_weight=decade_weight, cluster_weight=cluster_weight, device_weight=device_weight, loss_type=loss_type, loss_params=loss_params ) @classmethod def get_model_config(cls, model_name: str, multi_task: bool = False) -> Dict: """Get default configuration for a model""" if model_name not in TRAINING_CONFIGS: logger.warning(f"No default config for {model_name}, using base config") config = TRAINING_CONFIGS.get('resnet50', {}).copy() else: config = TRAINING_CONFIGS[model_name].copy() # Add multi-task specific configurations if multi_task: config.update({ 'multi_task': True, 'decade_weight': 1.0, 'cluster_weight': 1.0, 'multitask_hidden_dim': 512, 'multitask_dropout': 0.3 }) return config @classmethod def get_finetune_config(cls, model_name: str) -> Dict: """Get fine-tuning configuration for a model""" model_family = get_model_family(model_name) return FINETUNE_CONFIGS.get(model_family, {}).copy() @classmethod def list_available_models(cls) -> List[str]: """List all available model architectures""" return list(MODEL_REGISTRY.keys()) @classmethod def get_model_info(cls, model_name: str, multi_task: bool = False) -> Dict: """Get detailed information about a model""" if model_name not in MODEL_REGISTRY: raise ValueError(f"Unknown model: {model_name}") # Create a temporary model to get info num_classes = {'decade': 5, 'cluster': 10} if multi_task else 5 model = cls.create_model(model_name, num_classes=num_classes, pretrained=False, multi_task=multi_task) info = { 'name': model_name, 'timm_name': MODEL_REGISTRY[model_name], 'multi_task': multi_task, 'num_parameters': sum(p.numel() for p in model.parameters()), 'num_trainable_parameters': sum(p.numel() for p in model.parameters() if p.requires_grad), 'default_config': cls.get_model_config(model_name, multi_task=multi_task), 'finetune_config': cls.get_finetune_config(model_name), } # Clean up del model return info def create_optimizer(model: nn.Module, config: Dict) -> torch.optim.Optimizer: """ Create optimizer based on configuration Args: model: Model to optimize config: Configuration dictionary Returns: Optimizer instance """ optimizer_name = config.get('optimizer', 'adamw') learning_rate = config.get('learning_rate', 1e-3) weight_decay = config.get('weight_decay', 1e-4) if optimizer_name.lower() == 'adamw': optimizer = torch.optim.AdamW( model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(0.9, 0.999) ) elif optimizer_name.lower() == 'adam': optimizer = torch.optim.Adam( model.parameters(), lr=learning_rate, weight_decay=weight_decay ) elif optimizer_name.lower() == 'sgd': optimizer = torch.optim.SGD( model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=weight_decay ) else: raise ValueError(f"Unknown optimizer: {optimizer_name}") return optimizer def create_scheduler(optimizer: torch.optim.Optimizer, config: Dict) -> torch.optim.lr_scheduler._LRScheduler: """ Create learning rate scheduler Args: optimizer: Optimizer instance config: Configuration dictionary Returns: Scheduler instance """ scheduler_name = config.get('scheduler', 'cosine') epochs = config.get('epochs', 30) if scheduler_name == 'cosine': scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=epochs, eta_min=1e-6 ) elif scheduler_name == 'step': scheduler = torch.optim.lr_scheduler.StepLR( optimizer, step_size=10, gamma=0.1 ) elif scheduler_name == 'exponential': scheduler = torch.optim.lr_scheduler.ExponentialLR( optimizer, gamma=0.95 ) elif scheduler_name == 'reduce_on_plateau': scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', factor=0.5, patience=5, verbose=True ) else: raise ValueError(f"Unknown scheduler: {scheduler_name}") return scheduler def freeze_backbone(model: nn.Module, freeze_ratio: float = 0.5): """ Freeze early layers of the model Args: model: Model instance freeze_ratio: Ratio of layers to freeze (0.0 to 1.0) """ # Handle multi-task models - only freeze backbone if isinstance(model, MultiTaskModel): target_model = model.backbone logger.info("Freezing backbone layers in multi-task model") else: target_model = model # Get all named parameters all_params = list(target_model.named_parameters()) num_to_freeze = int(len(all_params) * freeze_ratio) # Freeze early layers for i, (name, param) in enumerate(all_params): if i < num_to_freeze: param.requires_grad = False logger.debug(f"Froze layer: {name}") else: param.requires_grad = True # Count parameters for the whole model trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) total = sum(p.numel() for p in model.parameters()) logger.info(f"Froze {num_to_freeze}/{len(all_params)} backbone layers") logger.info(f"Trainable parameters: {trainable:,} / {total:,} ({trainable / total * 100:.1f}%)") if __name__ == "__main__": # Test model creation print("Testing ModelFactory with Multi-task Support...\n") # List available models print("Available models:") for model_name in ModelFactory.list_available_models(): print(f" - {model_name}") # Test creating a single-task model print("\nCreating single-task EfficientNet-B2...") single_model = ModelFactory.create_model('efficientnet-b2', num_classes=5) # Test creating a multi-task model print("\nCreating multi-task EfficientNet-B2...") multi_model = ModelFactory.create_model( 'efficientnet-b2', num_classes={'decade': 5, 'cluster': 8}, multi_task=True, multitask_config={ 'hidden_dim': 512, 'dropout_rate': 0.3 } ) # Get model info single_info = ModelFactory.get_model_info('efficientnet-b2', multi_task=False) multi_info = ModelFactory.get_model_info('efficientnet-b2', multi_task=True) print(f"\nSingle-task model parameters: {single_info['num_parameters']:,}") print(f"Multi-task model parameters: {multi_info['num_parameters']:,}") # Test forward pass dummy_input = torch.randn(2, 3, 260, 260) print("\nTesting single-task forward pass...") single_output = single_model(dummy_input) print(f"Single-task output shape: {single_output.shape}") print("\nTesting multi-task forward pass...") multi_output = multi_model(dummy_input) print(f"Multi-task output shapes:") for task, output in multi_output.items(): print(f" {task}: {output.shape}") # Test multi-task loss print("\nTesting multi-task loss...") loss_fn = ModelFactory.create_multitask_loss(decade_weight=1.0, cluster_weight=0.8) targets = { 'decade': torch.randint(0, 5, (2,)), 'cluster': torch.randint(0, 8, (2,)) } losses = loss_fn(multi_output, targets) print(f"Loss components:") for loss_name, loss_value in losses.items(): print(f" {loss_name}: {loss_value.item():.4f}") print("\n✅ All tests passed!")