import torch import torch.nn.functional as F from torch.utils.data import DataLoader from transformers import Trainer from src.utils import batch_to_device from src.classifier_utils import HomogeneousBatchSampler # 手动实现 Focal Loss (带 alpha/gamma 控制) def sigmoid_focal_loss(inputs, targets, alpha: float = 0.25, gamma: float = 2.0, reduction: str = "mean"): """ Loss = -alpha * (1 - p)^gamma * log(p) """ p = torch.sigmoid(inputs) ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none") p_t = p * targets + (1 - p) * (1 - targets) loss = ce_loss * ((1 - p_t) ** gamma) if alpha >= 0: alpha_t = alpha * targets + (1 - alpha) * (1 - targets) loss = alpha_t * loss if reduction == "mean": return loss.mean() elif reduction == "sum": return loss.sum() return loss class EarlyExitTrainer(Trainer): def __init__(self, backbone_model, target_layer_idx, model_args, *args, **kwargs): self.max_length = kwargs.pop("max_length", 512) super().__init__(*args, **kwargs) self.backbone = backbone_model.to(self.args.device) self.backbone.eval() self.target_layer_idx = target_layer_idx self.model_args = model_args self._grad_check_done = False # ---------------- Dataloader ---------------- def get_train_dataloader(self) -> DataLoader: if self.train_dataset is None: raise ValueError("Trainer: training requires a train_dataset.") train_sampler = HomogeneousBatchSampler( self.train_dataset, batch_size=self._train_batch_size, drop_last=self.args.dataloader_drop_last, ) return DataLoader( self.train_dataset, batch_sampler=train_sampler, collate_fn=self.data_collator, num_workers=self.args.dataloader_num_workers, pin_memory=self.args.dataloader_pin_memory, ) # ---------------- Optimizer ---------------- def create_optimizer(self): if self.optimizer is None: print(f"\n[Debug Rank {self.args.local_rank}] Creating Optimizer...") decay_parameters = [] no_decay_parameters = [] trainable_count = 0 for name, param in self.model.named_parameters(): if not param.requires_grad: continue trainable_count += 1 if "bias" in name or "LayerNorm" in name or "BatchNorm" in name: no_decay_parameters.append(param) else: decay_parameters.append(param) print(f"[Debug] Found {trainable_count} trainable parameters.") self.optimizer = torch.optim.AdamW( [ {"params": decay_parameters, "weight_decay": self.args.weight_decay}, {"params": no_decay_parameters, "weight_decay": 0.0}, ], lr=self.args.learning_rate, eps=self.args.adam_epsilon, ) return self.optimizer # ---------------- Pooling ---------------- def _perform_pooling(self, hidden_state, attention_mask): """ 修复版 Pooling：增加索引边界检查，防止 CUDA Index Out of Bounds """ pooling_method = self.model_args.pooling batch_size, seq_len, _ = hidden_state.shape if pooling_method in ("last", "eos"): if attention_mask.shape[1] > seq_len: attention_mask = attention_mask[:, :seq_len] eos_indices = attention_mask.sum(dim=1) - 1 eos_indices = eos_indices.clamp(min=0, max=seq_len - 1) indices_expanded = eos_indices.unsqueeze(1).unsqueeze(2).expand(batch_size, 1, hidden_state.size(-1)) reps = torch.gather(hidden_state, 1, indices_expanded).squeeze(1) else: reps = hidden_state[:, -1, :] if self.model_args.normalize: reps = F.normalize(reps, p=2, dim=-1) return reps # ---------------- Loss entry ---------------- def compute_loss(self, model, inputs, return_outputs=False, **kwargs): loss = self._compute_early_exit_loss(model, inputs) return (loss, None) if return_outputs else loss # ---------------- Core loss (V5) ---------------- def _compute_early_exit_loss(self, model, inputs) -> torch.Tensor: self.backbone.eval() model.train() device = self.args.device qry_inputs, tgt_inputs = inputs qry_inputs = batch_to_device(qry_inputs, device) tgt_inputs = batch_to_device(tgt_inputs, device) # === 1. Backbone 特征提取 (无梯度) === with torch.no_grad(): # Target 侧 tgt_outputs = self.backbone.encoder( **tgt_inputs, return_dict=True, output_hidden_states=True ) tgt_hidden_mid = tgt_outputs.hidden_states[self.target_layer_idx] tgt_reps_mid = self._perform_pooling(tgt_hidden_mid, tgt_inputs["attention_mask"]) tgt_reps_last = self._perform_pooling(tgt_outputs.hidden_states[-1], tgt_inputs["attention_mask"]) # Query 侧 qry_outputs = self.backbone.encoder( **qry_inputs, return_dict=True, output_hidden_states=True ) q_hidden_mid = qry_outputs.hidden_states[self.target_layer_idx] q_hidden_last = qry_outputs.hidden_states[-1] qry_reps_mid = self._perform_pooling(q_hidden_mid, qry_inputs["attention_mask"]) qry_reps_last = self._perform_pooling(q_hidden_last, qry_inputs["attention_mask"]) batch_size = qry_reps_mid.size(0) backbone_ptr = self.backbone.module if hasattr(self.backbone, "module") else self.backbone temp = getattr(backbone_ptr, "temperature", 0.02) # 相似度计算 cos_mid = torch.matmul(qry_reps_mid, tgt_reps_mid.T) cos_last = torch.matmul(qry_reps_last, tgt_reps_last.T) scores_mid = cos_mid / temp probs_mid = torch.softmax(scores_mid, dim=1) # === 特征构造 (27维) === diag_cos = cos_mid.max(dim=1)[0] sorted_cos, _ = torch.sort(cos_mid, dim=1, descending=True) s2_cos = sorted_cos[:, 1] if sorted_cos.size(1) > 1 else sorted_cos[:, 0] margin_mid = diag_cos - s2_cos # 统计量 margin_mean = margin_mid.mean() margin_std = margin_mid.std(unbiased=False) + 1e-6 z_margin_mid = (margin_mid - margin_mean) / margin_std margin_median = margin_mid.median() mad = (margin_mid - margin_median).abs().median() + 1e-6 mad_margin_mid = (margin_mid - margin_median) / mad p1_mid = probs_mid.max(dim=1)[0] H_mid = -(probs_mid * torch.log(probs_mid + 1e-6)).sum(dim=1) gini_mid = 1.0 - (probs_mid ** 2).sum(dim=1) # TopK Stats TOPK = min(16, probs_mid.size(1)) topk_vals, _ = torch.topk(probs_mid, k=TOPK, dim=1) topk_mean = topk_vals.mean(dim=1) topk_std = topk_vals.std(dim=1, unbiased=False) topk_cv = topk_std / (topk_mean + 1e-6) centered = topk_vals - topk_mean.unsqueeze(1) var = (centered ** 2).mean(dim=1) + 1e-6 m4 = (centered ** 4).mean(dim=1) topk_kurt = m4 / (var ** 2) topk_med = topk_vals.median(dim=1).values row_mean_cos = cos_mid.mean(dim=1) row_med_cos = cos_mid.median(dim=1).values s1_over_mean = diag_cos - row_mean_cos s1_over_med = diag_cos - row_med_cos sorted_probs, _ = torch.sort(probs_mid, dim=1, descending=True) p1 = sorted_probs[:, 0] p2 = sorted_probs[:, 1] if sorted_probs.size(1) > 1 else sorted_probs[:, 0] shape_H = -(sorted_probs * torch.log(sorted_probs + 1e-6)).sum(dim=1) shape_gini = 1.0 - (sorted_probs ** 2).sum(dim=1) R = min(10, sorted_probs.size(1)) x = torch.arange(R, device=device, dtype=sorted_probs.dtype) x_centered = x - x.mean() denom = (x_centered ** 2).sum() y = torch.log(sorted_probs[:, :R] + 1e-6) slope = (x_centered.unsqueeze(0) * y).sum(dim=1) / denom row_mean_p = probs_mid.mean(dim=1) row_std_p = probs_mid.std(dim=1, unbiased=False) + 1e-6 z1 = (p1_mid - row_mean_p) / row_std_p center_p = probs_mid - row_mean_p.unsqueeze(1) m3 = (center_p ** 3).mean(dim=1) skew = m3 / (row_std_p ** 3 + 1e-6) s1_over_sk = p1_mid - skew TAIL_K = min(10, sorted_probs.size(1)) tail_mean = sorted_probs[:, -TAIL_K:].mean(dim=1) HEAD_K = min(5, sorted_probs.size(1)) head5_mean = sorted_probs[:, :HEAD_K].mean(dim=1) mask_ratio = torch.zeros_like(diag_cos) mask_len = torch.zeros_like(diag_cos) mask_runs = torch.zeros_like(diag_cos) scalar_inputs = torch.stack( [ diag_cos, s2_cos, margin_mid, z_margin_mid, mad_margin_mid, p1_mid, H_mid, gini_mid, topk_mean, topk_std, topk_cv, topk_kurt, topk_med, s1_over_mean, s1_over_med, p1, p2, shape_H, shape_gini, slope, z1, s1_over_sk, tail_mean, head5_mean, mask_ratio, mask_len, mask_runs, ], dim=1, ) # Modality Index modality_idx = torch.zeros(batch_size, dtype=torch.long, device=device) if "pixel_values" in qry_inputs and qry_inputs["pixel_values"] is not None: pv = qry_inputs["pixel_values"] if isinstance(pv, list): for i, item in enumerate(pv): if item is not None: modality_idx[i] = 1 elif isinstance(pv, torch.Tensor) and pv.numel() > 0: modality_idx.fill_(1) # Labels gt = torch.arange(batch_size, device=device) mid_top1 = cos_mid.argmax(dim=1) last_top1 = cos_last.argmax(dim=1) mid_hit = mid_top1.eq(gt) last_hit = last_top1.eq(gt) need_last = (~mid_hit) & last_hit labels = need_last.float().unsqueeze(1) both_correct = mid_hit & last_hit both_wrong = (~mid_hit) & (~last_hit) # ======================================================= # [V5关键修改] # 1. 缩进调整：以下逻辑在 no_grad 之外 # 2. 类型转换：强制转为 float32 (适应 V5 分类器的 TTA) # ======================================================= # 转换 scalar_inputs 和 qry_emb 为 float32 (分类器默认参数类型) scalar_inputs_f32 = scalar_inputs.float() qry_reps_mid_f32 = qry_reps_mid.float() # Forward logits = model(scalar_inputs_f32, modality_idx, qry_emb=qry_reps_mid_f32) # Loss: V5调整版 Focal Loss (alpha=0.80, gamma=3.0) # 重点关注正样本(Need Last)和难样本 loss = sigmoid_focal_loss(logits, labels, alpha=0.80, gamma=3.0, reduction="mean") pred_probs = torch.sigmoid(logits) # Logging if self.state.global_step < 10 and self.args.local_rank == 0: pos_ratio = labels.mean().item() neg_ratio = 1.0 - pos_ratio print(f"\n[Probe Step {self.state.global_step}] Loss: {loss.item():.4f}") print(f" - Pred Probs (need_last=1): mean={pred_probs.mean().item():.4f}, std={pred_probs.std().item():.4f}") print(f" - Labels: need_last={pos_ratio:.4f}, safe={neg_ratio:.4f}") print(f" - mid_hit: {mid_hit.float().mean().item():.2%}, last_hit: {last_hit.float().mean().item():.2%}") print(f" - both_correct: {both_correct.float().mean().item():.2%}, both_wrong: {both_wrong.float().mean().item():.2%}") return loss # ---------------- training_step ---------------- def training_step(self, model, inputs, num_items_in_batch=None) -> torch.Tensor: model.train() inputs = self._prepare_inputs(inputs) with self.compute_loss_context_manager(): loss = self.compute_loss(model, inputs) if self.args.n_gpu > 1: loss = loss.mean() self.accelerator.backward(loss) if not self._grad_check_done and self.args.local_rank == 0: print(f"\n[Gradient Check After Backward - Step {self.state.global_step}]") inner_model = model.module if hasattr(model, "module") else model has_grad = False total_grad_norm = 0.0 for name, param in inner_model.named_parameters(): if param.grad is not None: has_grad = True grad_norm = param.grad.norm().item() total_grad_norm += grad_norm ** 2 total_grad_norm = total_grad_norm ** 0.5 print(f" - Total Grad Norm: {total_grad_norm:.6f}") print(f" - Has Gradient: {has_grad}") if self.state.global_step >= 2: self._grad_check_done = True return loss.detach() / self.args.gradient_accumulation_steps