#!/usr/bin/env python3 """ 改进的多任务训练器（修正版） 修复/增强点： 1) Scheduler/ warmup 步数按 optimizer update steps（考虑 gradient_accumulation_steps） 2) eval_steps / save_steps 按 global_step(update step) 触发，而不是按 epoch 3) 支持最后不足 accumulation 的尾 batch 也执行一次 optimizer.step() 4) 支持 AMP（bf16 / fp16），fp16 使用 GradScaler 5) optimizer 仅包含 requires_grad=True 参数（适配 LoRA） 6) 训练/验证指标写入 jsonl，便于可视化与复现实验 7) 尝试启用 use_reentrant=False 的 gradient checkpointing（更推荐的实现） """ import os import math import json from dataclasses import asdict from datetime import datetime from pathlib import Path from contextlib import nullcontext 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 numpy as np try: import wandb except Exception: wandb = None from model_loader import ( load_model_and_processor, prepare_model_inputs ) from config import PretrainConfig class FocalLoss(nn.Module): """ Focal Loss for accident detection (保留，当前训练主损失仍使用 outputs.loss) """ def __init__(self, alpha=0.25, gamma=2.0): super().__init__() self.alpha = alpha self.gamma = gamma def forward(self, inputs, targets): 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: """ 多任务训练器 - 支持单帧任务 + 序列任务（取决于 batch 是否包含 "single_frame"/"sequence"） - 使用 LM loss (outputs.loss) 作为主要训练信号 """ def __init__( self, config: PretrainConfig, train_loader: DataLoader, val_loader: DataLoader, pretrained_lora_path: str | None = 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.global_step = 0 # optimizer update steps self.best_val_loss = float("inf") # 输出目录 self.output_dir = Path(self.config.training.output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) # 日志落盘（jsonl） self.train_log_path = self.output_dir / "train_metrics.jsonl" self.val_log_path = self.output_dir / "val_metrics.jsonl" # Loss functions（保留） self.ce_loss = nn.CrossEntropyLoss() self.focal_loss = FocalLoss(alpha=0.25, gamma=2.0) # Task weights（你的事故任务 + 兼容 bdd_* 默认 1.0） 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, # bdd_attributes / bdd_detection / bdd_drivable / bdd_risk 等，缺省走 1.0 } # Curriculum state（保留：注意你目前并未实际重建 dataloader） self.current_stage = 0 self.stage_epochs = [1, 2, 2] # AMP 配置 self.use_bf16 = bool(getattr(self.config.training, "bf16", False)) self.use_fp16 = bool(getattr(self.config.training, "fp16", False)) and (not self.use_bf16) self.autocast_dtype = torch.bfloat16 if self.use_bf16 else (torch.float16 if self.use_fp16 else None) self.scaler = torch.cuda.amp.GradScaler(enabled=self.use_fp16) # 初始化 self._init_model() self._init_optimizer() # wandb if self.config.training.use_wandb: if wandb is None: print("⚠️ wandb 未安装或不可用，但 use_wandb=True；将跳过 wandb 记录。") self.config.training.use_wandb = False else: wandb.init( project=self.config.training.wandb_project, name=self.config.training.wandb_run_name or f"{self.config.model.model_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}", config=asdict(self.config), ) # ------------------------- utils ------------------------- def _autocast_ctx(self): if self.autocast_dtype is None: return nullcontext() return torch.cuda.amp.autocast(dtype=self.autocast_dtype) def _write_jsonl(self, path: Path, record: dict): record = dict(record) record.setdefault("time", datetime.now().isoformat(timespec="seconds")) with open(path, "a", encoding="utf-8") as f: f.write(json.dumps(record, ensure_ascii=False) + "\n") def _move_batch_to_device(self, inputs: dict): """ 将 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 # ------------------------- init model/optim ------------------------- def _init_model(self): print("=" * 60) print("加载模型...") self.model, self.processor = load_model_and_processor(self.config.model) # 可选：Stage B 加载 Stage A 的 LoRA adapter 继续训练 if self.pretrained_lora_path: try: from peft import PeftModel self.model = PeftModel.from_pretrained( self.model, self.pretrained_lora_path, is_trainable=True ) print(f"✓ 已加载LoRA权重: {self.pretrained_lora_path}") except Exception as e: print(f"⚠️ 加载LoRA权重失败（将继续使用当前模型的LoRA状态）: {e}") 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 # 尽量启用更推荐的 checkpoint 实现（避免 use_reentrant 默认变化/限制） if hasattr(self.model, "gradient_checkpointing_enable"): try: self.model.gradient_checkpointing_enable( gradient_checkpointing_kwargs={"use_reentrant": False} ) except TypeError: # 旧版本 transformers/pytorch 不支持 kwargs try: self.model.gradient_checkpointing_enable() except Exception: pass try: if hasattr(self.model, "config"): self.model.config.use_cache = False except Exception: pass print(f"✓ 模型加载完成: {self.config.model.model_name}") print("=" * 60) def _init_optimizer(self): # 只优化 requires_grad=True 的参数（LoRA 正常情况下只有 adapter 参数是 True） trainable_params = [p for p in self.model.parameters() if p.requires_grad] if len(trainable_params) == 0: raise RuntimeError("没有可训练参数（requires_grad=True 为 0）。请检查 LoRA/冻结策略。") self.optimizer = torch.optim.AdamW( trainable_params, lr=self.config.training.learning_rate, weight_decay=self.config.training.weight_decay ) # Scheduler：关键修复 —— total steps 必须按 optimizer update steps 计算 grad_acc = max(1, int(self.config.training.gradient_accumulation_steps)) if len(self.train_loader) == 0: raise RuntimeError("train_loader 为空（len=0），请检查数据是否正确生成。") updates_per_epoch = math.ceil(len(self.train_loader) / grad_acc) num_training_steps = updates_per_epoch * int(self.config.training.num_epochs) num_warmup_steps = int(num_training_steps * float(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("✓ 优化器初始化完成") print(f" batches/epoch: {len(self.train_loader)}") print(f" grad_accum: {grad_acc}") print(f" updates/epoch: {updates_per_epoch}") print(f" total update steps: {num_training_steps}") print(f" warmup steps: {num_warmup_steps}") # ------------------------- label construction helpers ------------------------- 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_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 new_labels[i, :seq_len] = seq new_labels[i, :pl] = -100 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 与模型真实 input_ids 对齐（包含视觉 token） """ images = batch_data["images"] 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, 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, 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: fallback_inputs = self._concat_answers_to_prompt_inputs(prompt_inputs, labels_text) fallback_inputs = self._move_batch_to_device(fallback_inputs) return fallback_inputs, labels_text, task_types def prepare_sequence_inputs_and_labels(self, batch_data): """ 序列任务：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, 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) fallback_inputs = self._move_batch_to_device(fallback_inputs) return fallback_inputs, labels_text, task_types # ------------------------- loss / forward ------------------------- def compute_loss(self, batch): """ 计算 batch 总 loss（支持单帧 + 序列） """ total_loss = 0.0 task_losses = {} n_tasks = 0 with self._autocast_ctx(): if "single_frame" in batch: sf_data = batch["single_frame"] inputs, _, task_types = self.prepare_inputs_and_labels(sf_data) outputs = self.model(**inputs) loss = outputs.loss task_type = task_types[0] weight = float(self.task_weights.get(task_type, 1.0)) total_loss = total_loss + (loss * weight) task_losses[task_type] = float(loss.detach().item()) n_tasks += 1 if "sequence" in batch: seq_data = batch["sequence"] inputs, _, task_types = self.prepare_sequence_inputs_and_labels(seq_data) outputs = self.model(**inputs) loss = outputs.loss task_type = task_types[0] weight = float(self.task_weights.get(task_type, 1.0)) total_loss = total_loss + (loss * weight) task_losses[task_type] = float(loss.detach().item()) n_tasks += 1 if n_tasks > 0: total_loss = total_loss / n_tasks return total_loss, task_losses # ------------------------- eval / save ------------------------- @torch.no_grad() def evaluate(self, epoch: int): self.model.eval() val_loss_sum = 0.0 val_task_losses = {} for batch in tqdm(self.val_loader, desc="Validation"): loss, task_losses = self.compute_loss(batch) val_loss_sum += float(loss.detach().item()) for t, v in task_losses.items(): val_task_losses.setdefault(t, []).append(v) val_loss = val_loss_sum / max(1, len(self.val_loader)) avg_task_losses = {t: float(np.mean(vs)) for t, vs in val_task_losses.items()} record = {"step": self.global_step, "epoch": epoch, "val/loss": float(val_loss)} for t, v in avg_task_losses.items(): record[f"val/{t}"] = float(v) self._write_jsonl(self.val_log_path, record) if self.config.training.use_wandb and wandb is not None: wandb.log(record, step=self.global_step) return val_loss, avg_task_losses def _rotate_checkpoints(self): limit = int(getattr(self.config.training, "save_total_limit", 0) or 0) if limit <= 0: return ckpts = sorted( [p for p in self.output_dir.glob("checkpoint-*") if p.is_dir()], key=lambda p: int(p.name.split("-")[-1]) if p.name.split("-")[-1].isdigit() else 0 ) if len(ckpts) <= limit: return for p in ckpts[:-limit]: try: for sub in sorted(p.rglob("*"), reverse=True): if sub.is_file(): sub.unlink() elif sub.is_dir(): sub.rmdir() p.rmdir() except Exception: pass def save_checkpoint(self, tag: str, is_best: bool = False): if is_best: save_dir = self.output_dir / "best_model" else: save_dir = self.output_dir / tag save_dir.mkdir(parents=True, exist_ok=True) # 保存模型（LoRA/PEFT 会保存 adapter 权重） 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) # 保存 trainer state torch.save({ "epoch_tag": tag, "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}") if not is_best: self._rotate_checkpoints() # ------------------------- train loop ------------------------- def train(self): print("\n" + "=" * 60) print("开始训练") print("=" * 60) grad_acc = max(1, int(self.config.training.gradient_accumulation_steps)) logging_steps = max(1, int(self.config.training.logging_steps)) # 注意：这里的 eval_steps / save_steps 现在是按 global_step(update step) 生效 eval_steps = int(self.config.training.eval_steps) if self.config.training.eval_steps else 0 save_steps = int(self.config.training.save_steps) if self.config.training.save_steps else 0 total_epochs = int(self.config.training.num_epochs) # logging window stats window_task_sum = {} window_task_cnt = {} window_loss_sum = 0.0 window_updates = 0 for epoch in range(total_epochs): self.model.train() 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) self.optimizer.zero_grad(set_to_none=True) pbar = tqdm(self.train_loader, desc=f"Epoch {epoch+1}/{total_epochs}") for step, batch in enumerate(pbar): loss, task_losses = self.compute_loss(batch) # 梯度累积：缩放 loss loss_for_backward = loss / grad_acc if self.scaler.is_enabled(): self.scaler.scale(loss_for_backward).backward() else: loss_for_backward.backward() # accumulate task stats（记录未缩放 loss） for t, v in task_losses.items(): window_task_sum[t] = window_task_sum.get(t, 0.0) + float(v) window_task_cnt[t] = window_task_cnt.get(t, 0) + 1 # 是否该更新（包含尾 batch flush） do_update = ((step + 1) % grad_acc == 0) or ((step + 1) == len(self.train_loader)) if not do_update: continue # clip grad if self.scaler.is_enabled(): self.scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_( self.model.parameters(), float(self.config.training.max_grad_norm) ) # optimizer.step + scheduler.step（顺序很重要：先 optimizer 再 scheduler） if self.scaler.is_enabled(): self.scaler.step(self.optimizer) self.scaler.update() else: self.optimizer.step() self.scheduler.step() self.optimizer.zero_grad(set_to_none=True) # global step（update step） self.global_step += 1 # window logging window_updates += 1 window_loss_sum += float(loss.detach().item()) if self.global_step % logging_steps == 0: avg_loss = window_loss_sum / max(1, window_updates) log_dict = { "step": self.global_step, "epoch": epoch, "train/loss": float(avg_loss), "train/lr": float(self.scheduler.get_last_lr()[0]), } for t in window_task_sum: log_dict[f"train/{t}"] = float(window_task_sum[t] / max(1, window_task_cnt.get(t, 1))) # GPU memory（可选但很实用） if torch.cuda.is_available(): log_dict["train/gpu_mem_alloc_mb"] = float(torch.cuda.memory_allocated() / 1024 / 1024) log_dict["train/gpu_mem_reserved_mb"] = float(torch.cuda.memory_reserved() / 1024 / 1024) self._write_jsonl(self.train_log_path, log_dict) if self.config.training.use_wandb and wandb is not None: wandb.log(log_dict, step=self.global_step) pbar.set_postfix({ "loss": f"{avg_loss:.4f}", "lr": f"{log_dict['train/lr']:.2e}" }) # reset window window_task_sum, window_task_cnt = {}, {} window_loss_sum, window_updates = 0.0, 0 # save by steps if save_steps > 0 and (self.global_step % save_steps == 0): self.save_checkpoint(tag=f"checkpoint-{self.global_step}", is_best=False) # eval by steps（Stage A 验证集很大，eval_steps 不要设太小） if eval_steps > 0 and (self.global_step % eval_steps == 0): val_loss, val_task_losses = self.evaluate(epoch) print("\nValidation Results:") print(f" Overall Loss: {val_loss:.4f}") for t, v in val_task_losses.items(): print(f" {t}: {v:.4f}") if val_loss < self.best_val_loss: self.best_val_loss = val_loss self.save_checkpoint(tag="best_model", is_best=True) print(f"✓ 新的最佳模型! Val Loss: {val_loss:.4f}") # 每个 epoch 结束强制 eval（更稳妥） val_loss, val_task_losses = self.evaluate(epoch) print("\n[Epoch End] Validation Results:") print(f" Overall Loss: {val_loss:.4f}") for t, v in val_task_losses.items(): print(f" {t}: {v:.4f}") if val_loss < self.best_val_loss: self.best_val_loss = val_loss self.save_checkpoint(tag="best_model", is_best=True) print(f"✓ 新的最佳模型! Val Loss: {val_loss:.4f}") # Curriculum stage 更新（仍保留提示；如要真正生效应重建 dataloader） if self.current_stage < 3: if epoch + 1 == sum(self.stage_epochs[: self.current_stage + 1]): self.current_stage += 1 print(f"\n>>> Curriculum升级到 Stage {self.current_stage} <<<\n") # 最终保存 self.save_checkpoint(tag=f"checkpoint-{self.global_step}", is_best=False) print("\n" + "=" * 60) print("训练完成!") print(f"最佳验证Loss: {self.best_val_loss:.4f}") print(f"模型保存在: {self.output_dir}") print("=" * 60) if self.config.training.use_wandb and wandb is not None: wandb.finish()