import os from pathlib import Path from datetime import timedelta import time from datetime import datetime import wandb import torch from torch import nn from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter from accelerate import Accelerator from accelerate.utils import DistributedDataParallelKwargs, InitProcessGroupKwargs from transformer_maskgit.transformer_maskgit.optimizer import get_optimizer from scripts.scheduler import cosine_lr # === 导入两个 Dataset === from Rad_IR.radir.data.hmclip_mimic_box import HyCoClipMimicBoxDataset_JPG, collate_hmclip_box def cycle(dl): """让 DataLoader 无限循环""" while True: for data in dl: yield data class AverageMeter: """统计均值的工具类""" def __init__(self): self.reset() def reset(self): self.sum = 0 self.count = 0 self.val = 0 self.avg = 0 def update(self, val, n=1): """更新一次计数""" self.val = float(val) self.sum += float(val) * n self.count += n self.avg = self.sum / self.count if self.count > 0 else 0.0 class RadIRTrainer(nn.Module): """ RadIR 模型训练器： - 自动支持 Stage1（Box 监督） """ def __init__( self, RADIR: nn.Module, *, tokenizer, train_dataset, valid_dataset, num_train_steps, current_step=0, warmup_steps=1000, batch_size=32, lr=5e-5, wd=0.0, max_grad_norm=0.5, save_results_every=500, save_model_every=2000, results_folder="./outputs", num_workers=4, pin_memory=False, accelerate_kwargs=dict(), modal_embedding=False, debug=False, ): super().__init__() self.best_val_loss = float('inf') self.debug = debug self.modal_embedding = modal_embedding self.save_results_every = save_results_every self.save_model_every = save_model_every # ------------------------- # 判断当前是 Stage1 # ------------------------- self.is_stage1 = True # 只用 Stage1 print(f"\n[Trainer 初始化] 当前模式: Stage1（Simple）\n") # ------------------------- # 初始化 Accelerator（DDP + 混合精度） # ------------------------- ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) pg_kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=36000)) self.accelerator = Accelerator( mixed_precision="no", kwargs_handlers=[ddp_kwargs, pg_kwargs], **accelerate_kwargs, ) self.Rad_IR = RADIR self.tokenizer = tokenizer self.device = self.accelerator.device # ------------------------- # 基本训练参数 # ------------------------- self.num_train_steps = int(num_train_steps) self.current_step = int(current_step) self.batch_size = batch_size self.max_grad_norm = max_grad_norm # ------------------------- # 优化器 + 学习率调度 # ------------------------- all_params = list(self.Rad_IR.parameters()) init_lr = lr[0] if isinstance(lr, (list, tuple)) else float(lr) self.optim = get_optimizer(all_params, lr=init_lr, wd=wd) self.lr_scheduler = cosine_lr( self.optim, init_lr, warmup_steps, self.num_train_steps ) # ------------------------- # DataLoader（根据 stage 自动切换 collate） # ------------------------- collate_fn = collate_hmclip_box self.train_dl = DataLoader( train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=pin_memory, drop_last=True, collate_fn=collate_fn, ) self.valid_dl = DataLoader( valid_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=pin_memory, drop_last=False, collate_fn=collate_fn, ) # ------------------------- # Accelerator 包装模型、优化器、DataLoader # ------------------------- ( self.Rad_IR, self.optim, self.train_dl, self.valid_dl, ) = self.accelerator.prepare( self.Rad_IR, self.optim, self.train_dl, self.valid_dl, ) self.train_iter = cycle(self.train_dl) # ------------------------- # 日志 & 输出目录 # ------------------------- self.results_folder = Path(results_folder) self.results_folder.mkdir(parents=True, exist_ok=True) tb_path = self.results_folder / "tensorboard" tb_path.mkdir(parents=True, exist_ok=True) self.tb_writer = SummaryWriter(str(tb_path)) # ------------------------- # 统计指标 # ------------------------- self.uncon_loss_m = AverageMeter() # 总 loss self.it_ce_loss_m = AverageMeter() # image-text CE self.it_entailment_loss_m = AverageMeter() # image-text entailment self.ii_text_image_entailment_loss_m = AverageMeter() # image-image entailment print(f"[Trainer] 已准备完毕，总训练步数：{self.num_train_steps}") # ------------------------- # 时间统计 # ------------------------- self.train_start_time = None self.step_time_m = AverageMeter() # ======================================================= # 前向：根据 Stage 自动切换分支 # ======================================================= def _forward_batch(self, batch): imgs = batch["imgs"].to(self.device) caption_ids = batch["caption_ids"].to(self.device) token_type_ids = batch["token_type_ids"].to(self.device) attention_mask = batch["attention_mask"].to(self.device) phrase_text_mask = batch["phrase_text_mask"].to(self.device) phrases_bboxes_mask = batch["phrases_bboxes_mask"].to(self.device) text_batch = { "caption_ids": caption_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask, } loss, ce_loss, entail_loss, global_entail_loss = self.Rad_IR( text=text_batch, image=imgs, device=self.device, is_condition=False, return_loss=True, modal_embedding=self.modal_embedding, phrase_text_mask=phrase_text_mask, phrases_bboxes_mask=phrases_bboxes_mask, return_latents=True, ) return loss, ce_loss, entail_loss, global_entail_loss # ======================================================= # 主训练循环 # ======================================================= def train_model(self, resume_path=None): # ====================== # 先尝试 Resume 恢复训练 # ====================== if resume_path is not None: self.load_checkpoint(resume_path) self.Rad_IR.train() self.train_start_time = time.time() torch.autograd.set_detect_anomaly(True) for step in range(self.current_step, self.num_train_steps): self.current_step = step step_start_time = time.time() self.lr_scheduler.step() # ============================ # ✅ Stage 1 # ============================ with self.accelerator.accumulate(self.Rad_IR): batch = next(self.train_iter) loss, ce_loss, entail_loss, global_entail_loss = self._forward_batch(batch) #加检查1 for name, val in { "loss": loss, "ce": ce_loss, "entail": entail_loss, "global_entail": global_entail_loss, }.items(): if not torch.isfinite(val): print(f"[NaN LOSS] step={step}, {name}") self._dump_bad_batch(batch, step) raise RuntimeError("NaN in loss") #加检查2 self.optim.zero_grad() self.accelerator.backward(loss) #加检查1 for n, p in self.Rad_IR.named_parameters(): if p.grad is not None and not torch.isfinite(p.grad).all(): print(f"[NaN GRAD] step={step}, param={n}") raise RuntimeError("NaN in grad") #加检查2 # ============================ # ✅ 只有在 sync_gradients 才 step # ============================ if self.accelerator.sync_gradients: if self.max_grad_norm is not None: self.accelerator.clip_grad_norm_( self.Rad_IR.parameters(), self.max_grad_norm ) self.optim.step() self.optim.zero_grad(set_to_none=True) # ============================ # ✅ 时间统计 # ============================ if torch.cuda.is_available(): torch.cuda.synchronize() step_time = time.time() - step_start_time self.step_time_m.update(step_time) # ✅ loss 记录（已有的 AverageMeter） # ===== 从 GPU 拿回 4 个 loss ===== loss_val = float(loss.detach().cpu()) ce_loss_val = float(ce_loss.detach().cpu()) entail_loss_val = float(entail_loss.detach().cpu()) ii_entail_loss_val = float(global_entail_loss.detach().cpu()) # ===== 更新 4 个 meter ===== self.uncon_loss_m.update(loss_val, 1) self.it_ce_loss_m.update(ce_loss_val, 1) self.it_entailment_loss_m.update(entail_loss_val, 1) self.ii_text_image_entailment_loss_m.update(ii_entail_loss_val, 1) # ============================ # ✅ 日志 # ============================ if self.accelerator.is_main_process and (step + 1) % 10 == 0: eta_seconds = int(self.step_time_m.avg * (self.num_train_steps - step - 1)) eta_str = str(timedelta(seconds=eta_seconds)) print( f"[Step {step+1}/{self.num_train_steps}] " f"loss={self.uncon_loss_m.avg:.4f}, " f"ce={self.it_ce_loss_m.avg:.4f}, " f"entail={self.it_entailment_loss_m.avg:.4f}, " f"step_time={self.step_time_m.avg:.3f}s, ETA={eta_str}" ) wandb.log({ "train/loss": self.uncon_loss_m.avg, "train/ce_loss": self.it_ce_loss_m.avg, "train/entail_loss": self.it_entailment_loss_m.avg, "train/global_entail_loss": self.ii_text_image_entailment_loss_m.avg, "lr": self.optim.param_groups[0]["lr"], "step_time": self.step_time_m.avg, "step": step + 1 }, step=step + 1) # ============================ # ✅ 验证 # ============================ if (step + 1) % self.save_results_every == 0: val_loss = self.evaluate(step + 1) # ========= 更新 best ========= if self.accelerator.is_main_process: if val_loss < self.best_val_loss: self.best_val_loss = val_loss self.save_checkpoint(step + 1, is_best=True) # ============================ # ✅ 存模型 # ============================ if (step + 1) % self.save_model_every == 0: self.save_checkpoint(step + 1) # ✅ 最终评估 + 保存 self.evaluate(self.num_train_steps) self.save_checkpoint(self.num_train_steps, is_final=True) # ======================================================= # 验证集评估 # ======================================================= @torch.no_grad() def evaluate(self, step): self.Rad_IR.eval() eval_start_time = time.time() # ✅ 验证开始计时 loss_m = AverageMeter() ce_m = AverageMeter() entail_m = AverageMeter() g_entail_m = AverageMeter() for batch in self.valid_dl: loss, ce_loss, entail_loss, g_entail_loss = self._forward_batch(batch) bs = batch["imgs"].size(0) loss_m.update(loss.item(), bs) ce_m.update(ce_loss.item(), bs) entail_m.update(entail_loss.item(), bs) g_entail_m.update(g_entail_loss.item(), bs) if self.accelerator.is_main_process: print( f"[Eval step {step}] val_loss={loss_m.avg:.4f}, " f"val_ce={ce_m.avg:.4f}, val_entail={entail_m.avg:.4f}" ) self.tb_writer.add_scalar("val/loss", loss_m.avg, step) self.tb_writer.add_scalar("val/ce_loss", ce_m.avg, step) self.Rad_IR.train() eval_time = time.time() - eval_start_time eval_time_str = str(timedelta(seconds=int(eval_time))) if self.accelerator.is_main_process: print(f"[Eval Done] step={step} 耗时：{eval_time_str}") self.tb_writer.add_scalar("time/eval_time", eval_time, step) wandb.log({ "val/loss": loss_m.avg, "val/ce_loss": ce_m.avg, "val/entail_loss": entail_m.avg, "val/global_entail_loss": g_entail_m.avg, }, step=step) return loss_m.avg # ======================================================= # 模型保存 # ======================================================= def save_checkpoint(self, step, is_final=False, is_best=False): if not self.accelerator.is_main_process: return if is_final: ckpt_name = "radir_final.pt" elif is_best: ckpt_name = "radir_best.pt" else: ckpt_name = "radir_last.pt" # 覆盖最新 ckpt_path = self.results_folder / ckpt_name state = { "model": self.accelerator.get_state_dict(self.Rad_IR), "optimizer": self.optim.state_dict(), "scheduler": self.lr_scheduler.state_dict(), "step": step, "best_val_loss": self.best_val_loss, } torch.save(state, ckpt_path) print(f"[Checkpoint] 保存: {ckpt_path}") # ======================================================= # 模型恢复 # ======================================================= def load_checkpoint(self, ckpt_path): if not os.path.exists(ckpt_path): raise FileNotFoundError(f"Checkpoint 不存在: {ckpt_path}") print(f"[Resume] 从 {ckpt_path} 恢复训练") ckpt = torch.load(ckpt_path, map_location="cpu") # 1️⃣ 恢复模型 self.accelerator.unwrap_model(self.Rad_IR).load_state_dict(ckpt["model"]) # 2️⃣ 恢复优化器 self.optim.load_state_dict(ckpt["optimizer"]) # 3️⃣ 恢复 scheduler（如果存在） if "scheduler" in ckpt: self.lr_scheduler.load_state_dict(ckpt["scheduler"]) print("[Resume] Scheduler 已恢复") else: print("[Resume] ⚠️ 未发现 scheduler，使用新的 scheduler") # 4️⃣ 恢复训练进度 self.current_step = ckpt["step"] # 5️⃣ 恢复 best_val_loss self.best_val_loss = ckpt.get("best_val_loss", float("inf")) print("resume lr =", self.optim.param_groups[0]["lr"]) print( f"[Resume] 恢复完成：step = {self.current_step}, " f"best_val_loss = {self.best_val_loss:.4f}" )