scripts/RadIR_Trainer_CXR_CT.py

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}"
        )