#!/usr/bin/env python3 """ 改进的多任务训练器 融合最新VLM领域适应研究: 1. Curriculum Learning 2. Dynamic Task Weighting 3. Contrastive Learning for accident detection """ import os import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from transformers import get_linear_schedule_with_warmup from tqdm import tqdm import wandb from pathlib import Path import json from datetime import datetime import numpy as np from model_loader import ( load_model_and_processor, prepare_model_inputs ) from config import PretrainConfig class FocalLoss(nn.Module): """ Focal Loss for accident detection 处理正负样本不平衡问题 """ def __init__(self, alpha=0.25, gamma=2.0): super().__init__() self.alpha = alpha self.gamma = gamma def forward(self, inputs, targets): """ Args: inputs: predicted logits [B, num_classes] targets: target labels [B] """ ce_loss = F.cross_entropy(inputs, targets, reduction='none') p_t = torch.exp(-ce_loss) focal_loss = self.alpha * (1 - p_t) ** self.gamma * ce_loss return focal_loss.mean() class MultiTaskTrainer: """ 多任务训练器 支持Curriculum Learning和动态任务权重 """ def __init__(self, config: PretrainConfig, train_loader: DataLoader, val_loader: DataLoader, pretrained_lora_path: str = None): self.config = config self.train_loader = train_loader self.val_loader = val_loader self.pretrained_lora_path = pretrained_lora_path # 设备 self.device = torch.device(config.training.device) print(f"使用设备: {self.device}") # 加载模型 self._init_model() # 初始化优化器和scheduler self._init_optimizer() # Loss functions self.ce_loss = nn.CrossEntropyLoss() self.focal_loss = FocalLoss(alpha=0.25, gamma=2.0) # Task weights (可以动态调整) self.task_weights = { "scene_understanding": self.config.data.task1_weight, "binary_detection": self.config.data.task2_weight, "accident_description": self.config.data.task3_weight, "sequence_prediction": self.config.data.task3_weight, # 使用task3权重 # Stage A (BDD100K) tasks "bdd_attributes": 1.0, "bdd_detection": 1.0, "bdd_drivable": 1.0, "bdd_risk": 1.0 } # Curriculum learning state self.current_stage = 0 # 0=easy, 1=medium, 2=hard, 3=all self.stage_epochs = [1, 2, 2] # 每个stage的epoch数 # 训练统计 self.global_step = 0 self.best_val_loss = float('inf') self.train_losses = [] # wandb if config.training.use_wandb: wandb.init( project=config.training.wandb_project, name=config.training.wandb_run_name or f"{config.model.model_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}", config=vars(config) ) def _init_model(self): """初始化模型""" print("=" * 60) print("加载模型...") self.model, self.processor = load_model_and_processor(self.config.model) # 如果提供了预训练LoRA权重，加载它 if self.pretrained_lora_path: print(f"\n加载预训练LoRA权重: {self.pretrained_lora_path}") # 检查路径 lora_path = Path(self.pretrained_lora_path) if not lora_path.exists(): print(f"❌ LoRA权重不存在: {lora_path}") raise FileNotFoundError(f"LoRA权重不存在: {lora_path}") # 加载LoRA权重 try: from peft import PeftModel self.model = PeftModel.from_pretrained( self.model, self.pretrained_lora_path, is_trainable=True # 继续训练 ) print("✓ 预训练LoRA权重加载成功") print("✓ LoRA权重设置为可训练") except Exception as e: print(f"❌ LoRA权重加载失败: {e}") raise self.model.to(self.device) # 确保tokenizer有pad_token if self.processor.tokenizer.pad_token is None: self.processor.tokenizer.pad_token = self.processor.tokenizer.eos_token self.processor.tokenizer.pad_token_id = self.processor.tokenizer.eos_token_id print(f"✓ 模型加载完成: {self.config.model.model_name}") if self.pretrained_lora_path: print(f"✓ 从预训练LoRA继续训练: {self.pretrained_lora_path}") print("=" * 60) def _init_optimizer(self): """初始化优化器和scheduler""" # 优化器 self.optimizer = torch.optim.AdamW( self.model.parameters(), lr=self.config.training.learning_rate, weight_decay=self.config.training.weight_decay ) # Scheduler num_training_steps = len(self.train_loader) * self.config.training.num_epochs num_warmup_steps = int(num_training_steps * self.config.training.warmup_ratio) self.scheduler = get_linear_schedule_with_warmup( self.optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps ) print(f"✓ 优化器初始化完成") print(f" 总步数: {num_training_steps}") print(f" warmup: {num_warmup_steps}") def _move_batch_to_device(self, inputs): """将processor输出的BatchFeature/dict移动到device，并对浮点张量对齐到模型dtype。""" moved = {} for k, v in inputs.items(): if torch.is_tensor(v): if v.is_floating_point(): moved[k] = v.to(self.device, dtype=self.model.dtype) else: moved[k] = v.to(self.device) else: moved[k] = v return moved def _concat_answers_to_prompt_inputs(self, prompt_inputs, labels_text): """Fallback：在无法用processor重新编码full_texts时，把答案tokens拼接到prompt input_ids后面并构造labels。""" tokenizer = self.processor.tokenizer pad_id = tokenizer.pad_token_id eos_id = tokenizer.eos_token_id input_ids = prompt_inputs["input_ids"] attention_mask = prompt_inputs["attention_mask"] if input_ids.dim() != 2 or attention_mask.dim() != 2: raise ValueError("prompt_inputs 必须包含二维的 input_ids 和 attention_mask") B, L = input_ids.shape prompt_lens = attention_mask.sum(dim=1).tolist() # 每条样本的 answer token（不加special tokens），末尾追加 eos answer_ids_list = [ tokenizer.encode(ans, add_special_tokens=False) + ([eos_id] if eos_id is not None else []) for ans in labels_text ] max_full_len = max(int(pl) + len(ans_ids) for pl, ans_ids in zip(prompt_lens, answer_ids_list)) new_input_ids = input_ids.new_full((B, max_full_len), pad_id) new_attention_mask = attention_mask.new_zeros((B, max_full_len)) new_labels = input_ids.new_full((B, max_full_len), -100) for i, (pl, ans_ids) in enumerate(zip(prompt_lens, answer_ids_list)): pl = int(pl) ans_tensor = torch.tensor(ans_ids, device=input_ids.device, dtype=input_ids.dtype) seq = torch.cat([input_ids[i, :pl], ans_tensor], dim=0) seq_len = seq.size(0) new_input_ids[i, :seq_len] = seq new_attention_mask[i, :seq_len] = 1 # labels 等于 full input_ids，但 prompt 部分 mask new_labels[i, :seq_len] = seq new_labels[i, :pl] = -100 # 组装新的inputs：保留与视觉相关的张量/元信息；丢弃任何基于旧seq_len的二维张量（例如position_ids） out = {} for k, v in prompt_inputs.items(): if k in ("input_ids", "attention_mask", "labels"): continue if torch.is_tensor(v) and v.dim() == 2 and v.shape[0] == B and v.shape[1] == L: continue out[k] = v out["input_ids"] = new_input_ids out["attention_mask"] = new_attention_mask out["labels"] = new_labels return out def prepare_inputs_and_labels(self, batch_data): """ 准备单帧任务的模型输入和labels（labels与模型真实input_ids严格对齐，包含视觉token长度） """ images = batch_data["images"] user_prompts = batch_data["user_prompts"] labels_text = batch_data["labels"] task_types = batch_data["task"] # 先走一次prepare_model_inputs，复用其中生成的prompt_texts（chat template + 视觉占位符等） prompt_inputs = prepare_model_inputs( self.processor, self.config.model.model_type, images, user_prompts, self.device ) prompt_texts = prompt_inputs.pop("__prompt_texts__", None) if prompt_texts is None: raise RuntimeError("prepare_model_inputs 未返回 __prompt_texts__，无法构造对齐labels") full_texts = [ prompt + answer + self.processor.tokenizer.eos_token for prompt, answer in zip(prompt_texts, labels_text) ] # 优先用processor(文本+图像)得到包含视觉token展开后的真实input_ids try: prompt_encodings = self.processor( text=prompt_texts, images=images, return_tensors="pt", padding=True, truncation=True, max_length=512 ) full_inputs = self.processor( text=full_texts, images=images, return_tensors="pt", padding=True, truncation=True, max_length=512 ) labels = full_inputs["input_ids"].clone() for i in range(labels.size(0)): prompt_len = int(prompt_encodings["attention_mask"][i].sum().item()) labels[i, :prompt_len] = -100 labels[full_inputs["attention_mask"] == 0] = -100 full_inputs["labels"] = labels full_inputs = self._move_batch_to_device(full_inputs) return full_inputs, labels_text, task_types except Exception as e: # 如果processor无法处理输入（通常发生在多图/序列的特殊格式），回退到“在prompt input_ids后拼接答案token”的方式 fallback_inputs = self._concat_answers_to_prompt_inputs(prompt_inputs, labels_text) return fallback_inputs, labels_text, task_types def prepare_sequence_inputs_and_labels(self, batch_data): """ 准备序列任务的模型输入和labels（与单帧同逻辑，但images是List[List[PIL]]或等价格式） """ images_list = batch_data["image_sequences"] user_prompts = batch_data["user_prompts"] labels_text = batch_data["labels"] task_types = batch_data["task"] prompt_inputs = prepare_model_inputs( self.processor, self.config.model.model_type, images_list, # List of List of images user_prompts, self.device ) prompt_texts = prompt_inputs.pop("__prompt_texts__", None) if prompt_texts is None: raise RuntimeError("prepare_model_inputs 未返回 __prompt_texts__，无法构造对齐labels") full_texts = [ prompt + answer + self.processor.tokenizer.eos_token for prompt, answer in zip(prompt_texts, labels_text) ] try: prompt_encodings = self.processor( text=prompt_texts, images=images_list, return_tensors="pt", padding=True, truncation=True, max_length=512 ) full_inputs = self.processor( text=full_texts, images=images_list, return_tensors="pt", padding=True, truncation=True, max_length=512 ) labels = full_inputs["input_ids"].clone() for i in range(labels.size(0)): prompt_len = int(prompt_encodings["attention_mask"][i].sum().item()) labels[i, :prompt_len] = -100 labels[full_inputs["attention_mask"] == 0] = -100 full_inputs["labels"] = labels full_inputs = self._move_batch_to_device(full_inputs) return full_inputs, labels_text, task_types except Exception: fallback_inputs = self._concat_answers_to_prompt_inputs(prompt_inputs, labels_text) return fallback_inputs, labels_text, task_types def compute_loss(self, batch): """ 计算batch的总loss 支持单帧和序列任务 """ total_loss = 0 task_losses = {} n_tasks = 0 # 处理单帧任务 if "single_frame" in batch: sf_data = batch["single_frame"] # 准备输入和标签 inputs, labels_text, task_types = self.prepare_inputs_and_labels(sf_data) # Forward pass outputs = self.model(**inputs) loss = outputs.loss # 根据任务类型加权 task_type = task_types[0] # batch中同一任务 weighted_loss = loss * self.task_weights.get(task_type, 1.0) total_loss += weighted_loss task_losses[task_type] = loss.item() n_tasks += 1 # 处理序列任务 if "sequence" in batch: seq_data = batch["sequence"] # 准备输入和labels（labels与input_ids对齐，包含视觉token长度） inputs, labels_text, task_types = self.prepare_sequence_inputs_and_labels(seq_data) # Forward pass outputs = self.model(**inputs) loss = outputs.loss # 根据任务类型加权 task_type = task_types[0] # batch中同一任务 weighted_loss = loss * self.task_weights.get(task_type, 1.0) total_loss += weighted_loss task_losses[task_type] = loss.item() n_tasks += 1 # 平均 if n_tasks > 0: total_loss = total_loss / n_tasks return total_loss, task_losses def train_epoch(self, epoch): """训练一个epoch""" self.model.train() epoch_loss = 0 epoch_task_losses = {} pbar = tqdm(self.train_loader, desc=f"Epoch {epoch+1}/{self.config.training.num_epochs}") for step, batch in enumerate(pbar): # 计算loss loss, task_losses = self.compute_loss(batch) # Backward loss = loss / self.config.training.gradient_accumulation_steps loss.backward() # 累积task losses for task, task_loss in task_losses.items(): if task not in epoch_task_losses: epoch_task_losses[task] = [] epoch_task_losses[task].append(task_loss) # 梯度累积 if (step + 1) % self.config.training.gradient_accumulation_steps == 0: # 梯度裁剪 torch.nn.utils.clip_grad_norm_( self.model.parameters(), self.config.training.max_grad_norm ) # 更新 self.optimizer.step() self.scheduler.step() self.optimizer.zero_grad() self.global_step += 1 # 记录 if self.global_step % self.config.training.logging_steps == 0: avg_task_losses = { task: np.mean(losses) for task, losses in epoch_task_losses.items() } log_dict = { "train/loss": loss.item() * self.config.training.gradient_accumulation_steps, "train/lr": self.scheduler.get_last_lr()[0], "train/step": self.global_step, "train/epoch": epoch } for task, avg_loss in avg_task_losses.items(): log_dict[f"train/{task}"] = avg_loss if self.config.training.use_wandb: wandb.log(log_dict, step=self.global_step) pbar.set_postfix({ "loss": f"{loss.item() * self.config.training.gradient_accumulation_steps:.4f}", "lr": f"{self.scheduler.get_last_lr()[0]:.2e}" }) epoch_loss += loss.item() return epoch_loss / len(self.train_loader) @torch.no_grad() def evaluate(self): """验证""" self.model.eval() val_loss = 0 val_task_losses = {} for batch in tqdm(self.val_loader, desc="Validation"): loss, task_losses = self.compute_loss(batch) val_loss += loss.item() for task, task_loss in task_losses.items(): if task not in val_task_losses: val_task_losses[task] = [] val_task_losses[task].append(task_loss) val_loss /= len(self.val_loader) avg_task_losses = { task: np.mean(losses) for task, losses in val_task_losses.items() } return val_loss, avg_task_losses def save_checkpoint(self, epoch, is_best=False): """保存checkpoint""" # 保存目录 if is_best: save_dir = Path(self.config.training.output_dir) / "best_model" else: save_dir = Path(self.config.training.output_dir) / f"checkpoint-{epoch}" save_dir.mkdir(parents=True, exist_ok=True) # 保存模型（LoRA或完整模型） if hasattr(self.model, "save_pretrained"): self.model.save_pretrained(save_dir) else: torch.save(self.model.state_dict(), save_dir / "pytorch_model.bin") # 保存processor self.processor.save_pretrained(save_dir) # 保存训练状态 torch.save({ "epoch": epoch, "global_step": self.global_step, "optimizer_state_dict": self.optimizer.state_dict(), "scheduler_state_dict": self.scheduler.state_dict(), "best_val_loss": self.best_val_loss, }, save_dir / "trainer_state.pt") print(f"✓ 保存checkpoint: {save_dir}") def train(self): """主训练循环 - 支持Curriculum Learning""" print("\n" + "=" * 60) print("开始训练") print("=" * 60) # Curriculum Learning: 逐stage训练 # Stage 0: easy (1 epoch) # Stage 1: medium (2 epochs) # Stage 2: hard (2 epochs) # Stage 3: all (remaining epochs) total_epochs = self.config.training.num_epochs for epoch in range(total_epochs): print(f"\n{'='*60}") print(f"Epoch {epoch+1}/{total_epochs}") print(f"Curriculum Stage: {self.current_stage} ({['easy', 'medium', 'hard', 'all'][min(self.current_stage, 3)]})") print("=" * 60) # 训练 train_loss = self.train_epoch(epoch) # 验证 if (epoch + 1) % self.config.training.eval_steps == 0 or epoch == total_epochs - 1: val_loss, val_task_losses = self.evaluate() print(f"\nValidation Results:") print(f" Overall Loss: {val_loss:.4f}") for task, loss in val_task_losses.items(): print(f" {task}: {loss:.4f}") # Wandb if self.config.training.use_wandb: log_dict = {"val/loss": val_loss, "val/epoch": epoch} for task, loss in val_task_losses.items(): log_dict[f"val/{task}"] = loss wandb.log(log_dict, step=self.global_step) # 保存最佳模型 if val_loss < self.best_val_loss: self.best_val_loss = val_loss self.save_checkpoint(epoch, is_best=True) print(f"✓ 新的最佳模型! Val Loss: {val_loss:.4f}") # 定期保存 if (epoch + 1) % self.config.training.save_steps == 0: self.save_checkpoint(epoch) # Curriculum stage 更新 if self.current_stage < 3: # 根据stage_epochs更新stage if epoch + 1 == sum(self.stage_epochs[:self.current_stage+1]): self.current_stage += 1 print(f"\n>>> Curriculum升级到 Stage {self.current_stage} <<<\n") # 注意：实际应用中需要重新创建dataloader # 这里为简化，保持当前loader # 最终保存 self.save_checkpoint(total_epochs - 1) print("\n" + "=" * 60) print("训练完成!") print(f"最佳验证Loss: {self.best_val_loss:.4f}") print(f"模型保存在: {self.config.training.output_dir}") print("=" * 60) if self.config.training.use_wandb: wandb.finish()