training/pretrain/trainer_v2_modified.py

#!/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()