Spaces:

readylab
/

VizRef

Sleeping

VizRef / src /models /model_factory.py

Add model and inference code

18a82fb about 1 year ago

21.4 kB

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!")

	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!")