"""Token-based distillation training task for models with distillation heads.""" import logging from typing import Dict, Optional, Union import torch import torch.nn as nn import torch.nn.functional as F from timm.models import create_model from timm.utils import unwrap_model from .task import TrainingTask _logger = logging.getLogger(__name__) class TokenDistillationTeacher(nn.Module): """Wrapper for a teacher model used in token-based distillation. Creates and manages a pre-trained teacher model for token distillation, handling model creation and normalization differences between teacher and student. Can be created from: - A model name string (creates the model internally) - An existing nn.Module (wraps it with the necessary interface) Args: model_name_or_module: Either a model name string or an nn.Module num_classes: Number of output classes (required if model_name_or_module is a string) in_chans: Number of input channels (used if model_name_or_module is a string) pretrained_path: Optional path to pretrained weights (used if model_name_or_module is a string) device: Device to place the model on dtype: Model dtype (uses float32 if None) """ def __init__( self, model_name_or_module: Union[str, nn.Module], num_classes: Optional[int] = None, in_chans: int = 3, pretrained_path: Optional[str] = None, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, ): super().__init__() if isinstance(model_name_or_module, str): _logger.info(f"Creating token distillation teacher model: '{model_name_or_module}'") pretrained_kwargs = {'pretrained': True} if pretrained_path: pretrained_kwargs['pretrained_cfg_overlay'] = dict( file=pretrained_path, num_classes=num_classes, ) model = create_model( model_name=model_name_or_module, num_classes=num_classes, in_chans=in_chans, device=device, dtype=dtype, **pretrained_kwargs, ) elif isinstance(model_name_or_module, nn.Module): model = model_name_or_module else: raise TypeError( f"model_name_or_module must be a string or nn.Module, got {type(model_name_or_module).__name__}" ) model.eval() self.model = model # Get normalization values from pretrained_cfg if available model_unwrapped = unwrap_model(model) if hasattr(model_unwrapped, 'pretrained_cfg'): mean = model_unwrapped.pretrained_cfg.get('mean', (0.485, 0.456, 0.406)) std = model_unwrapped.pretrained_cfg.get('std', (0.229, 0.224, 0.225)) else: mean = (0.485, 0.456, 0.406) std = (0.229, 0.224, 0.225) mean_kd = torch.tensor(mean, device=device, dtype=dtype).view(1, -1, 1, 1) std_kd = torch.tensor(std, device=device, dtype=dtype).view(1, -1, 1, 1) self.register_buffer('mean_kd', mean_kd, persistent=False) self.register_buffer('std_kd', std_kd, persistent=False) def forward(self, input: torch.Tensor) -> torch.Tensor: """Forward pass through teacher model. Args: input: Input tensor (should already be normalized for teacher) Returns: Teacher logits """ return self.model(input) def normalize_input( self, input: torch.Tensor, student_mean: Optional[torch.Tensor] = None, student_std: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Normalize input to match teacher's expected normalization. Args: input: Input tensor (already normalized for student) student_mean: Student normalization mean buffer [1, 3, 1, 1] student_std: Student normalization std buffer [1, 3, 1, 1] Returns: Input tensor normalized for the teacher model """ if student_mean is None or student_std is None: return input if torch.equal(student_mean, self.mean_kd) and torch.equal(student_std, self.std_kd): return input return (input * student_std + student_mean - self.mean_kd) / self.std_kd class TokenDistillationTask(TrainingTask): """Token-based distillation task for models with distillation heads. For models like DeiT that have a dedicated distillation token/head that returns a tuple (main_logits, dist_logits) when distilled_training is enabled. The main head is trained against ground truth labels while the distillation head matches teacher outputs. Supports two distillation modes: - 'soft': KL divergence with temperature scaling (default) - 'hard': Cross-entropy with teacher's hard predictions (argmax) Loss weighting supports two modes: 1. Independent weights: loss = task_loss_weight * task_loss + distill_loss_weight * distill_loss 2. Complementary mode: loss = task_loss_weight * task_loss + (1 - task_loss_weight) * distill_loss (used when only task_loss_weight is specified) Args: student_model: Student model with set_distilled_training() method teacher_model: Teacher model - can be a model name string, nn.Module, or TokenDistillationTeacher criterion: Task loss function for main head (default: CrossEntropyLoss) teacher_pretrained_path: Path to teacher pretrained weights (used when teacher_model is a string) distill_type: 'soft' for KL-div or 'hard' for CE with teacher argmax distill_loss_weight: Weight for distillation loss task_loss_weight: Weight for task loss temperature: Softmax temperature for soft distillation (ignored for hard) device: Device for task tensors/buffers dtype: Dtype for task tensors/buffers verbose: Enable info logging Example: >>> # With model name string (num_classes/in_chans inferred from student) >>> task = TokenDistillationTask( ... student_model=model, teacher_model='deit_base_patch16_224', ... criterion=nn.CrossEntropyLoss(), ... distill_type='soft', temperature=3.0, task_loss_weight=0.5, ... device=torch.device('cuda'), ... ) >>> # With raw model >>> task = TokenDistillationTask( ... student_model=model, teacher_model=my_teacher_model, ... criterion=nn.CrossEntropyLoss(), ... distill_type='hard', task_loss_weight=0.5, ... ) """ def __init__( self, student_model: nn.Module, teacher_model: Union[str, nn.Module, TokenDistillationTeacher], criterion: Optional[nn.Module] = None, teacher_pretrained_path: Optional[str] = None, distill_type: str = 'soft', distill_loss_weight: Optional[float] = None, task_loss_weight: Optional[float] = None, temperature: float = 1.0, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, verbose: bool = True, ): super().__init__(device=device, dtype=dtype, verbose=verbose) # Validate model has set_distilled_training method student_unwrapped = unwrap_model(student_model) if not hasattr(student_unwrapped, 'set_distilled_training'): raise ValueError( f"Model {student_unwrapped.__class__.__name__} does not have 'set_distilled_training' method. " "TokenDistillationTask requires a model with a distillation head (e.g., DeiT distilled variants)." ) # Enable distilled training mode student_unwrapped.set_distilled_training(True) # Handle different teacher input types if isinstance(teacher_model, TokenDistillationTeacher): teacher = teacher_model elif isinstance(teacher_model, str) or isinstance(teacher_model, nn.Module): # Get num_classes and in_chans from student num_classes = student_unwrapped.num_classes in_chans = student_unwrapped.in_chans teacher = TokenDistillationTeacher( model_name_or_module=teacher_model, num_classes=num_classes, in_chans=in_chans, pretrained_path=teacher_pretrained_path, device=self.device, dtype=self.dtype, ) else: raise TypeError( f"teacher_model must be a model name string, nn.Module, or TokenDistillationTeacher, " f"got {type(teacher_model).__name__}" ) self.student = student_model self.teacher = teacher self.criterion = criterion if criterion is not None else nn.CrossEntropyLoss() self.distill_type = distill_type self.temperature = temperature if distill_type not in ('soft', 'hard'): raise ValueError(f"Unsupported distill_type '{distill_type}'. Must be 'soft' or 'hard'.") # Register student normalization values as non-persistent buffers student_mean = torch.tensor( student_unwrapped.pretrained_cfg['mean'], device=self.device, dtype=self.dtype, ).view(1, -1, 1, 1) student_std = torch.tensor( student_unwrapped.pretrained_cfg['std'], device=self.device, dtype=self.dtype, ).view(1, -1, 1, 1) self.register_buffer('student_mean', student_mean, persistent=False) self.register_buffer('student_std', student_std, persistent=False) # Determine weighting mode if distill_loss_weight is not None: # Mode 1: distill_weight specified - independent weights (task defaults to 1.0 if not set) self.distill_loss_weight = distill_loss_weight self.task_loss_weight = task_loss_weight if task_loss_weight is not None else 1.0 if self.verbose: _logger.info( f"TokenDistillationTask: Independent weights - " f"task_weight={self.task_loss_weight}, distill_weight={distill_loss_weight}" ) elif task_loss_weight is not None: # Mode 2: only task_weight specified - complementary mode (distill = 1 - task) self.task_loss_weight = task_loss_weight self.distill_loss_weight = 1.0 - task_loss_weight if self.verbose: _logger.info( f"TokenDistillationTask: Complementary mode - " f"task_weight={task_loss_weight}, distill_weight={self.distill_loss_weight}" ) else: # Mode 3: neither specified - equal weights (both 1.0) self.distill_loss_weight = 1.0 self.task_loss_weight = 1.0 if self.verbose: _logger.info( f"TokenDistillationTask: Default equal weights - " f"task_weight={self.task_loss_weight}, distill_weight={self.distill_loss_weight}" ) if self.verbose: _logger.info( f"TokenDistillationTask: distill_type={distill_type}, temperature={temperature}" ) def prepare_distributed( self, device_ids: Optional[list] = None, **ddp_kwargs ) -> 'TokenDistillationTask': """Prepare task for distributed training. Wraps the student model in DistributedDataParallel (DDP) while leaving the frozen teacher model unwrapped. Args: device_ids: List of device IDs for DDP (e.g., [local_rank]) **ddp_kwargs: Additional arguments passed to DistributedDataParallel Returns: self (for method chaining) """ from torch.nn.parallel import DistributedDataParallel as DDP for param in self.teacher.parameters(): param.requires_grad = False self.student = DDP(self.student, device_ids=device_ids, **ddp_kwargs) return self def forward( self, input: torch.Tensor, target: torch.Tensor, ) -> Dict[str, torch.Tensor]: """Forward pass with token distillation. Args: input: Input tensor [B, C, H, W] target: Target labels [B] Returns: Dictionary containing: - 'loss': Combined training loss (task + distillation) - 'output': Main head logits (for metrics) - 'task_loss': Classification loss component - 'distill_loss': Distillation loss component """ # Student forward pass - returns tuple (main_logits, dist_logits) student_output = self.student(input) main_logits, dist_logits = student_output # Compute task loss on main head task_loss = self.criterion(main_logits, target) # Teacher forward pass (no gradient) with torch.no_grad(): input_kd = self.teacher.normalize_input(input, self.student_mean, self.student_std) teacher_logits = self.teacher(input_kd.detach()) # Compute distillation loss on distillation head if self.distill_type == 'soft': prob_s = F.log_softmax(dist_logits / self.temperature, dim=-1) prob_t = F.log_softmax(teacher_logits / self.temperature, dim=-1) distill_loss = F.kl_div(prob_s, prob_t, reduction='batchmean', log_target=True) * (self.temperature ** 2) else: teacher_hard = teacher_logits.argmax(dim=-1) distill_loss = F.cross_entropy(dist_logits, teacher_hard) total_loss = self.task_loss_weight * task_loss + self.distill_loss_weight * distill_loss return { 'loss': total_loss, 'output': main_logits, 'task_loss': task_loss, 'distill_loss': distill_loss, }