train_rsb.py

import os

import argparse
from tensorboardX import SummaryWriter
from torch.optim.lr_scheduler import StepLR

from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions.normal import Normal
from torch.distributions.uniform import Uniform
from torch.optim import Adam
from torch.utils.data import DataLoader
from tqdm import tqdm
import matplotlib.pyplot as plt

from models.kan_resdiff_sb import Unet

from dataloaders import *
from metrics import *
from validate import validate

from models.pretrain.ProbUNet import ProbabilisticUNet, ProbUNetConfig
from models.pretrain.TRDataset import TRCacheByIdDataset
import yaml
from models.path_sb.diffusion import KAN_path, PathUNet
from models.path_sb.rsb import bridge_loss, rsb_diffuse
from datetime import datetime
from checkpoint_utils import load_model_state
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
def make_timestamp_folder(save_path, time_format="%Y-%m-%d_%H-%M-%S"):
    """
    在 save_path 下创建一个以当前时间命名的子文件夹，并返回其路径。

    Args:
        save_path (str): 父目录
        time_format (str): 时间格式（可自定义）

    Returns:
        str: 创建好的子目录路径
    """
    timestamp = datetime.now().strftime(time_format)
    final_path = os.path.join(save_path, timestamp)
    return final_path

def set_stage2_model_trainability(model, train_trnet):
    """Freeze the diffusion backbone; optionally enable TRNet parameters."""
    for name, p in model.named_parameters():
        if 'tr_encoder' in name:
            if train_trnet:
                p.requires_grad = True
            else:
                p.requires_grad = False
        else:
            p.requires_grad = False

def make_trnet_optimizer(model):
    """Return only TRNet parameters for the Stage 2 optimizer."""
    set_stage2_model_trainability(model, train_trnet=True)
    trnet_params = [p for p in model.parameters() if p.requires_grad]
    if len(trnet_params) == 0:
        raise ValueError("No trainable TRNet parameters found for Stage 2.")
    return trnet_params

def set_seed(seed):
    import random
    import numpy as np
    import torch
    import os

    random.seed(seed)
    np.random.seed(seed)

    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    
    torch.set_float32_matmul_precision('high')

    os.environ["PYTHONHASHSEED"] = str(seed)

    print(f"[INFO] Global seed set to {seed}")
    
def Dice_batch(pred, target):
    smooth = 1e-8
    tp = (pred * target).sum(dim=[1,2,3])
    denom = pred.sum(dim=[1,2,3]) + target.sum(dim=[1,2,3])
    return (2 * tp + smooth) / (denom + smooth)

def IoU_batch(pred, target):
    smooth = 1e-8
    tp = (pred * target).sum(dim=[1,2,3])
    fp = ((1 - target) * pred).sum(dim=[1,2,3])
    fn = (target * (1 - pred)).sum(dim=[1,2,3])
    return (tp + smooth) / (tp + fp + fn + smooth)

def save_topk_model(
    best_models,
    topk,
    current_metrics,   # dict: epoch, GED, MDM, IoU
    checkpoint,
    save_path
):
    """
    Maintain top-K models based on GED (lower is better)
    """
    epoch = current_metrics["epoch"]
    GED   = current_metrics["GED"]

    filename = f"Best@Epoch:{current_metrics['epoch']:04d}_MDM:{current_metrics['MDM']:.4f}_GED{current_metrics['GED']:.4f}_IoU:{current_metrics['IoU']:.4f}.pth"
    ckpt_path = os.path.join(save_path, filename)

    if len(best_models) < topk:
        torch.save(checkpoint, ckpt_path)
        best_models.append({**current_metrics, "path": ckpt_path})

    else:
        worst = max(best_models, key=lambda x: x["GED"])

        if GED < worst["GED"]:
            if os.path.exists(worst["path"]):
                os.remove(worst["path"])

            best_models.remove(worst)

            torch.save(checkpoint, ckpt_path)
            best_models.append({**current_metrics, "path": ckpt_path})

    best_models.sort(key=lambda x: x["GED"])

def make_ddpm_schedule(T=1000, beta_start=1e-4, beta_end=2e-2, device="cuda"):
    betas = torch.linspace(beta_start, beta_end, T, device=device)          # (T,)
    alphas = 1.0 - betas                                                   # (T,)
    alphas_bar = torch.cumprod(alphas, dim=0)                               # (T,)
    sqrt_ab = torch.sqrt(alphas_bar)                                        # (T,)
    sqrt_1mab = torch.sqrt(1.0 - alphas_bar)                                # (T,)
    return betas, alphas, alphas_bar, sqrt_ab, sqrt_1mab

@torch.no_grad()
def single_t_sample_ab(
    model,
    mask_gt,      # (N,1,H,W)
    img_cond,     # (N,1,H,W)
    tr_cond,      # (N,1,H,W)
    t_int,        # int 或 (N,)
    sqrt_ab,
    sqrt_1mab,
):
    """
    单次 t 扩散 + 反解 a,b
    """

    device = mask_gt.device
    N = mask_gt.shape[0]
    T = sqrt_ab.shape[0]

    if isinstance(t_int, int):
        t_batch = torch.full((N,), t_int, device=device, dtype=torch.long)
    else:
        t_batch = t_int

    a_true = sqrt_ab[t_batch].view(-1,1,1,1)
    b_true = sqrt_1mab[t_batch].view(-1,1,1,1)

    eps = torch.randn_like(mask_gt)

    x_t = a_true * mask_gt + b_true * eps

    t_norm = (t_batch.float() / (T - 1))

    eps_pred = model(x_t, t_norm, img_cond, tr_cond).float()

    x0_hat = (x_t - b_true * eps_pred) / (a_true + 1e-8)

    a_model, b_model = solve_a_b(
        x0=mask_gt,
        eps=eps_pred.detach(),
        x_pred=x_t
    )

    return {
        "a_true": a_true.squeeze(),
        "b_true": b_true.squeeze(),
        "a_model":  a_model,
        "b_model":  b_model,
        "x_t":    x_t,
        "x0_hat": x0_hat,
        "eps": eps,
        "eps_pred": eps_pred.detach(),
    }

def solve_a_b(x0, eps, x_pred, eps_stability=1e-8):
    """
    x0, eps, x_pred: (B,1,H,W)
    返回:
        a, b: (B,)
    """

    B = x0.shape[0]

    x0_flat = x0.view(B, -1)
    eps_flat = eps.view(B, -1)
    x_flat  = x_pred.view(B, -1)

    S00 = (x0_flat * x0_flat).sum(dim=1)
    S11 = (eps_flat * eps_flat).sum(dim=1)
    S01 = (x0_flat * eps_flat).sum(dim=1)

    S0x = (x0_flat * x_flat).sum(dim=1)
    S1x = (eps_flat * x_flat).sum(dim=1)

    Delta = S00 * S11 - S01 * S01 + eps_stability

    a = (S11 * S0x - S01 * S1x) / Delta
    b = (S00 * S1x - S01 * S0x) / Delta

    return a, b

def loss_fn_ddpm(
    model,
    mask_gt,     # (N,1,H,W)  真实mask
    t_int,       # (N,)
    img_cond,    # (N,1,H,W)  原图
    tr_cond,     # (N,1,H,W)  TR (可选条件)
    sqrt_ab,
    sqrt_1mab,
):
    eps = torch.randn_like(mask_gt)

    a = sqrt_ab[t_int].view(-1,1,1,1)
    b = sqrt_1mab[t_int].view(-1,1,1,1)

    x_t = a * mask_gt + b * eps

    t = (t_int.float() / (sqrt_ab.shape[0] - 1)).to(x_t.device)

    eps_pred = model(x_t, t, img_cond, tr_cond).float()

    loss_eps = (eps_pred - eps).pow(2).mean(dim=[1,2,3])

    x0_hat = (x_t - b * eps_pred) / (a + 1e-8)

    y_hat = x0_hat.clamp(0,1)
    y_bin = (y_hat >= 0.5).float()
    y_gt  = (mask_gt >= 0.5).float()

    dice = Dice_batch(y_bin, y_gt)
    iou  = IoU_batch(y_bin, y_gt)

    return loss_eps, dice, iou

def loss_fn_rsb_ddpm(
    model,
    mask_gt,       # (N,1,H,W)
    t_int,         # (N,)
    img_cond,      # (N,1,H,W)
    tr_cond,       # (N,1,H,W) or None
    sqrt_ab,
    sqrt_1mab,
    paths,         # dict with fa, fb, ba, bb
    use_rsb=True,
    T=1000,
):
    """
    Epsilon loss on either DDPM states or learned forward-path states.

    When use_rsb=True, x_t is built from fa/fb and x0 reconstruction uses ba/bb.
    TR conditioning is handled inside the model's TRNet branch.
    """
    eps = torch.randn_like(mask_gt)
    t_norm = (t_int.float() / (T - 1))

    if use_rsb:
        with torch.no_grad():
            a_f_pred = paths['fa'](eps, t_norm, img_cond).view(-1)
            b_f_pred = paths['fb'](eps, t_norm, img_cond).view(-1)
        x_t = rsb_diffuse(mask_gt, eps, a_f_pred, b_f_pred)
    else:
        a = sqrt_ab[t_int].view(-1, 1, 1, 1)
        b = sqrt_1mab[t_int].view(-1, 1, 1, 1)
        x_t = a * mask_gt + b * eps

    t = t_norm.to(x_t.device)

    tr_cond_scaled = tr_cond

    eps_pred = model(x_t, t, img_cond, tr_cond_scaled).float()
    loss_eps = (eps_pred - eps).pow(2).mean(dim=[1, 2, 3])

    if use_rsb:
        with torch.no_grad():
            a_b = paths['ba'](eps_pred, t_norm, img_cond).view(-1, 1, 1, 1)
            b_b = paths['bb'](eps_pred, t_norm, img_cond).view(-1, 1, 1, 1)
        x0_hat = (x_t - b_b * eps_pred) / (a_b + 1e-8)
    else:
        a = sqrt_ab[t_int].view(-1, 1, 1, 1)
        b = sqrt_1mab[t_int].view(-1, 1, 1, 1)
        x0_hat = (x_t - b * eps_pred) / (a + 1e-8)

    y_hat = x0_hat.clamp(0, 1)
    y_bin = (y_hat >= 0.5).float()
    y_gt = (mask_gt >= 0.5).float()

    dice = Dice_batch(y_bin, y_gt)
    iou = IoU_batch(y_bin, y_gt)

    return loss_eps, dice, iou

def train(
    paths,
    tr_dataset,
    tr_dataset_val,
    model,
    dataloader,
    optimizer,
    device,
    num_epochs,
    save_path,
    val_dataloader,
    scheduler,
    task_name,
    rsb_enabled=False,
    lambda_bridge=0.1,
    lambda_ddpm_anchor=1.0,
    lambda_alpha=0.01,
):
    """Train path models, then optionally fine-tune TRNet on path-defined states.

    With rsb_enabled=False, fa/fb and ba/bb are fitted around the fixed DDPM
    schedule. With rsb_enabled=True, the first loop alternates forward-path and
    backward-path updates, then the second loop freezes the backbone and trains
    TRNet on learned forward-path states.
    """
    TOPK = 3
    best_models = []  

    T = 1000
    _, _, _, sqrt_ab, sqrt_1mab = make_ddpm_schedule(T=T, device=device)

    save_path = os.path.join(save_path, task_name)
    save_path = make_timestamp_folder(save_path)

    tb_log_dir = os.path.join(save_path, "tb_logs")
    os.makedirs(tb_log_dir, exist_ok=True)
    writer = SummaryWriter(log_dir=tb_log_dir)

    model._stage(0)
    model.eval()
    best_backward_loss = 999
    best_bridge_loss = 999
    best_combined_metric = 999

    for epoch in range(num_epochs):

        is_odd_epoch = (epoch % 2 == 1)  # 奇数轮更新 forward path；偶数轮更新 backward path

        if is_odd_epoch:
            print(f"{task_name}; Epochs:{epoch+1}/{num_epochs} [Forward path] ... ")
        else:
            print(f"{task_name}; Epochs:{epoch+1}/{num_epochs} [Backward path] ... ")

        total_loss_forward = 0
        total_loss_backward = 0
        total_loss_bridge = 0

        pbar = tqdm(dataloader, ncols=100, leave=False)
        for images, masks, _, case_ids in pbar:
            optimizer.zero_grad()

            images = images.to(device)
            batch_size = images.shape[0]

            m1 = masks[0].to(device)
            m2 = masks[1].to(device)
            m3 = masks[2].to(device)
            m4 = masks[3].to(device)

            all_masks = torch.cat([m1, m2, m3, m4], dim=0)
            all_images = images.repeat(4, 1, 1, 1)

            t_int = torch.randint(low=0, high=T, size=(batch_size,), device=device, dtype=torch.long)
            t_batch = t_int.repeat(4)

            result = single_t_sample_ab(
                model=model,
                mask_gt=all_masks,
                t_int=t_batch,
                img_cond=all_images,
                tr_cond=None,
                sqrt_ab=sqrt_ab,
                sqrt_1mab=sqrt_1mab,
            )

            t_norm = (t_batch.float() / (T - 1))

            a_f = paths['fa'](result['eps'], t_norm, all_images).view(-1)
            b_f = paths['fb'](result['eps'], t_norm, all_images).view(-1)

            a_model = result['a_model'].detach()
            b_model = result['b_model'].detach()

            a_b = paths['ba'](result['eps_pred'], t_norm, all_images).view(-1)
            b_b = paths['bb'](result['eps_pred'], t_norm, all_images).view(-1)

            if rsb_enabled:

                a_diff = sqrt_ab[t_batch]
                b_diff = sqrt_1mab[t_batch]

                if is_odd_epoch:
                    loss_bridge_f = bridge_loss(a_f, b_f, a_b.detach(), b_b.detach())
                    loss_fa_base = ((a_f - a_diff) ** 2).mean()
                    loss_fb_base = ((b_f - b_diff) ** 2).mean()
                    loss_forward = (
                        lambda_bridge * loss_bridge_f +
                        lambda_ddpm_anchor * (loss_fa_base + loss_fb_base)
                    )
                    loss_backward = None  # 本轮不更新反向路径
                    paths["ob"].zero_grad()  # 清空 ob 梯度（防止残留）
                else:
                    loss_ba = ((a_b - a_model) ** 2).mean()
                    loss_bb = ((b_b - b_model) ** 2).mean()
                    loss_bridge_b = bridge_loss(a_f.detach(), b_f.detach(), a_b, b_b)
                    loss_backward = (
                        (loss_ba + loss_bb) +
                        lambda_bridge * loss_bridge_b
                    )
                    loss_forward = None  # 本轮不更新前向路径
                    paths["of"].zero_grad()  # 清空 of 梯度（防止残留）

            else:
                a_diff = sqrt_ab[t_batch]
                b_diff = sqrt_1mab[t_batch]

                loss_fa = ((a_f - a_diff) ** 2).mean()
                loss_fb = ((b_f - b_diff) ** 2).mean()
                loss_forward = loss_fa + loss_fb

                loss_ba = ((a_b - a_model) ** 2).mean()
                loss_bb = ((b_b - b_model) ** 2).mean()
                loss_backward = loss_ba + loss_bb

            if rsb_enabled:
                if is_odd_epoch:
                    paths["of"].zero_grad()
                    loss_forward.backward()
                    paths["of"].step()
                else:
                    paths["ob"].zero_grad()
                    loss_backward.backward()
                    paths["ob"].step()

                if is_odd_epoch:
                    total_loss_bridge += loss_bridge_f.item()
            else:
                paths["of"].zero_grad()
                loss_forward.backward()
                paths["of"].step()

                paths["ob"].zero_grad()
                loss_backward.backward()
                paths["ob"].step()

            if rsb_enabled:
                if is_odd_epoch:
                    total_loss_forward += loss_forward.item()
                else:
                    total_loss_backward += loss_backward.item()
            else:
                total_loss_forward += loss_forward.item()
                total_loss_backward += loss_backward.item()

            postfix = {}
            if rsb_enabled:
                if is_odd_epoch:
                    postfix["Fwd"] = f"{loss_forward.item():.4f}"
                else:
                    postfix["Bwd"] = f"{loss_backward.item():.4f}"
                if is_odd_epoch:
                    postfix["bridge"] = f"{loss_bridge_f.item():.4f}"
            else:
                postfix["loss(b)"] = f"{loss_backward:.4f}"
                postfix["loss(f)"] = f"{loss_forward:.4f}"
            pbar.set_postfix(postfix)

        os.makedirs(save_path, exist_ok=True)
        num_batches = len(dataloader)
        epoch_idx = epoch + 1

        if rsb_enabled:
            if is_odd_epoch:
                paths["osf"].step()
            else:
                paths["osb"].step()
        else:
            paths["osf"].step()
            paths["osb"].step()

        avg_loss_forward = total_loss_forward / num_batches if not rsb_enabled or is_odd_epoch else 0
        avg_loss_backward = total_loss_backward / num_batches if not rsb_enabled or not is_odd_epoch else 0

        if rsb_enabled:
            avg_loss_bridge = total_loss_bridge / num_batches if is_odd_epoch else 0

            if is_odd_epoch:
                print(f"[Epoch {epoch_idx}] Forward Step: Fwd={avg_loss_forward:.6f} Bridge={avg_loss_bridge:.6f}")
            else:
                print(f"[Epoch {epoch_idx}] Backward Step: Bwd={avg_loss_backward:.6f}")

            if is_odd_epoch and avg_loss_bridge < best_bridge_loss:
                best_bridge_loss = avg_loss_bridge
            if not is_odd_epoch and avg_loss_backward < best_backward_loss:
                best_backward_loss = avg_loss_backward

            if not is_odd_epoch:  # 偶数轮结束后评估
                combined_metric = avg_loss_backward + best_bridge_loss
                if combined_metric < best_combined_metric:
                    best_combined_metric = combined_metric
                    best_path = os.path.join(save_path, "paths_rsb.pth")
                    torch.save({
                        "epoch": epoch,
                        "model": model.state_dict(),
                        "paths": get_paths_state_rsb(paths),
                        "best_backward_loss": best_backward_loss,
                        "best_bridge_loss": best_bridge_loss,
                    }, best_path)
                    print(f"New best RSB model saved (bwd={avg_loss_backward:.6f}, bridge={best_bridge_loss:.6f}, combined={combined_metric:.6f})")
        else:
            print(f"[Epoch {epoch_idx}] Avg Forward Loss: {avg_loss_forward:.6f}")
            print(f"[Epoch {epoch_idx}] Avg Backward Loss: {avg_loss_backward:.6f}")

            if avg_loss_backward < best_backward_loss:
                best_backward_loss = avg_loss_backward
                best_path = os.path.join(save_path, "paths.pth")
                torch.save({
                    "epoch": epoch,
                    "model": model.state_dict(),
                    "paths": get_paths_state(paths),
                    "best_backward_loss": best_backward_loss,
                }, best_path)
                print(f"New best backward model saved (loss={best_backward_loss:.6f})")

    if not rsb_enabled:
        return
    for epoch in range(num_epochs):

        model._stage(1)
        set_stage2_model_trainability(model, train_trnet=True)
        if rsb_enabled:
            for k in ["fa", "fb", "ba", "bb"]:
                paths[k].eval()
                for p in paths[k].parameters():
                    p.requires_grad = False
        print(f"{task_name}; Epochs:{epoch+1}/{num_epochs} [Stage 2: TRNet] ... ")

        total_loss = 0
        total_dice = 0
        total_iou = 0
        pbar = tqdm(dataloader, ncols=100, leave=False)
        for images, masks, _, case_ids in pbar:
            optimizer.zero_grad()

            images = images.to(device)
            batch_size = images.shape[0]

            m1 = masks[0].to(device)
            m2 = masks[1].to(device)
            m3 = masks[2].to(device)
            m4 = masks[3].to(device)

            all_masks = torch.cat([m1, m2, m3, m4], dim=0)

            all_images = images.repeat(4, 1, 1, 1)

            t_int = torch.randint(low=0, high=T, size=(batch_size,), device=device, dtype=torch.long)
            t_int = t_int.repeat(4)   # (4B,)

            TR_list = [tr_dataset.get_by_id(cid) for cid in case_ids]  # each (1,H,W)
            TR = torch.stack(TR_list, dim=0).to(device, non_blocking=True)  # (B,1,H,W)

            TR_rep = TR.repeat(4, 1, 1, 1)   # (4B,1,H,W)

            if rsb_enabled:
                loss_all, dice_all, iou_all = loss_fn_rsb_ddpm(
                    model=model,
                    mask_gt=all_masks,
                    t_int=t_int,
                    img_cond=all_images,
                    tr_cond=TR_rep,
                    sqrt_ab=sqrt_ab,
                    sqrt_1mab=sqrt_1mab,
                    paths=paths,
                    use_rsb=True,
                    T=T,
                )
            else:
                loss_all, dice_all, iou_all = loss_fn_ddpm(
                    model=model,
                    mask_gt=all_masks,
                    t_int=t_int,
                    img_cond=all_images,
                    tr_cond=TR_rep,
                    sqrt_ab=sqrt_ab,
                    sqrt_1mab=sqrt_1mab,
                )

            loss_all = loss_all.view(4, batch_size)
            dice_all = dice_all.view(4, batch_size)
            iou_all  = iou_all.view(4, batch_size)

            loss = loss_all.mean()
            dice = dice_all.mean()
            iou  = iou_all.mean()

            back_loss = loss + 0.0001 * (1 - dice)

            if rsb_enabled and model.alpha is not None:
                t_norm = (t_int.float() / (T - 1))
                alpha_target = 1.0 - t_norm
                loss_alpha = lambda_alpha * ((model.alpha - alpha_target) ** 2).mean()
                back_loss = back_loss + loss_alpha

            back_loss.backward()
            optimizer.step()

            total_loss += loss
            total_dice += dice
            total_iou  += iou

            pbar.set_postfix({
                "loss": f"{loss.item():.4f}",
                "dice": f"{dice.item():.4f}",
                "iou":  f"{iou.item():.4f}",
            })

        scheduler.step()
        epoch_idx = epoch + 1  # 训练内的 epoch 编号，从 1 开始

        total_loss = total_loss / len(dataloader)
        total_dice = total_dice / len(dataloader)
        total_iou = total_iou / len(dataloader)
        print(f'Epoch:{epoch_idx}, Loss={total_loss}, Dice={total_dice}, IoU={total_iou}')
        

        validate_mode = "rsb" if rsb_enabled else "standard"
        validate_paths = paths if rsb_enabled else None
        _MDM, _GED, _IoU = validate(
            tr_dataset_val, model, val_dataloader, device,
            sampling_times=24, num_samples=4, print_full=False, tr_enabled=True,
            paths=validate_paths, sampling_mode=validate_mode,
        )
        f_MDM, f_GED, f_IoU = validate(
            tr_dataset_val, model, val_dataloader, device,
            sampling_times=24, num_samples=4, print_full=False, tr_enabled=False,
            paths=validate_paths, sampling_mode=validate_mode,
        )
        metric_str = f'Epoch:{epoch_idx:04d}-MDM:{_MDM:.4f}-GED:{_GED:.4f}-IoU:{_IoU:.4f}'
        metric_str_f = f'Epoch:{epoch_idx:04d}-MDM:{f_MDM:.4f}-GED:{f_GED:.4f}-IoU:{f_IoU:.4f}'
        d_MDM = _MDM - f_MDM
        d_GED = _GED - f_GED
        d_IoU = _IoU - f_IoU
        metric_str_diff = (
            f'Epoch:{epoch_idx:04d}-'
            f'ΔMDM:{d_MDM:+.4f}-'
            f'ΔGED:{d_GED:+.4f}-'
            f'ΔIoU:{d_IoU:+.4f}'
        )
        print(f"valid: {metric_str}")
        print(f"w/otr: {metric_str_f}")
        print(f"vdiff: {metric_str_diff}")

        writer.add_scalar("train/Loss", total_loss.item(), epoch_idx)
        writer.add_scalars(
            "train/metrics",
            {
                "Dice": total_dice.item(),
                "IoU":  total_iou.item(),
            },
            epoch_idx
        )
        writer.add_scalars(
            "val/diff",
            {
                "ΔMDM":     d_MDM,
                "ΔGED":     d_GED,
                "ΔIoU":     d_IoU,
            },
            epoch_idx
        )

        writer.add_scalars(
            "val/metrics",
            {
                "MDM":     _MDM,
                "GED":     _GED,
                "IoU":     _IoU,
            },
            epoch_idx
        )
        writer.add_scalars(
            "val/metrics_origin",
            {
                "MDM":     f_MDM,
                "GED":     f_GED,
                "IoU":     f_IoU,
            },
            epoch_idx
        )

        save_epoch_list = list(range(50, num_epochs + 1, 50))
        current_metrics = {
            "epoch": epoch_idx,
            "MDM":  _MDM,
            "GED":  _GED,
            "IoU":  _IoU,
        }

        checkpoint = {
            "epoch": epoch_idx,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "scheduler_state_dict": scheduler.state_dict(),
            "paths": get_paths_state_rsb(paths) if rsb_enabled else None,
            "loss": loss,
            "GED": _GED,
        }

        save_topk_model(
            best_models=best_models,
            topk=TOPK,
            current_metrics=current_metrics,
            checkpoint=checkpoint,
            save_path=save_path
        )

        if epoch_idx in save_epoch_list:
            print(f"Saving epoch model: {metric_str}")
            os.makedirs(save_path, exist_ok=True)
            torch.save(checkpoint, os.path.join(save_path, f"{metric_str}.pth"))

    print("\n==== Top-3 Best Models (by GED) ====")
    for i, m in enumerate(best_models):
        print(
            f"[{i+1}] Epoch {m['epoch']:04d} | "
            f"MDM={m['MDM']:.4f} | GED={m['GED']:.4f} | IoU={m['IoU']:.4f}"
        )

    writer.close()

def get_paths_state(paths):
    return {
        "fa": paths["fa"].state_dict(),
        "fb": paths["fb"].state_dict(),
        "ba": paths["ba"].state_dict(),
        "bb": paths["bb"].state_dict(),
        "optimizer_f": paths["of"].state_dict(),
        "optimizer_b": paths["ob"].state_dict(),
        "scheduler_f": paths["osf"].state_dict(),
        "scheduler_b": paths["osb"].state_dict(),
    }
def load_paths_state(paths, ckpt_dict, device=None):
    """
    Load saved states back into paths dict.

    Args:
        paths: 当前已构建好的 paths 字典
        ckpt_dict: torch.load 得到的字典（即 get_paths_state 保存的内容）
        device: 可选，若指定则强制 optimizer state 迁移到该 device
    """

    paths["fa"].load_state_dict(ckpt_dict["fa"])
    paths["fb"].load_state_dict(ckpt_dict["fb"])
    paths["ba"].load_state_dict(ckpt_dict["ba"])
    paths["bb"].load_state_dict(ckpt_dict["bb"])

    paths["of"].load_state_dict(ckpt_dict["optimizer_f"])
    paths["ob"].load_state_dict(ckpt_dict["optimizer_b"])

    paths["osf"].load_state_dict(ckpt_dict["scheduler_f"])
    paths["osb"].load_state_dict(ckpt_dict["scheduler_b"])

    if device is not None:
        for opt_key in ["of", "ob"]:
            for state in paths[opt_key].state.values():
                for k, v in state.items():
                    if isinstance(v, torch.Tensor):
                        state[k] = v.to(device)

def get_paths_state_rsb(paths):
    return {
        "fa": paths["fa"].state_dict(),
        "fb": paths["fb"].state_dict(),
        "ba": paths["ba"].state_dict(),
        "bb": paths["bb"].state_dict(),
        "optimizer_f": paths["of"].state_dict(),
        "optimizer_b": paths["ob"].state_dict(),
        "scheduler_f": paths["osf"].state_dict(),
        "scheduler_b": paths["osb"].state_dict(),
    }

def load_paths_state_rsb(paths, ckpt_dict, device=None):
    paths["fa"].load_state_dict(ckpt_dict["fa"])
    paths["fb"].load_state_dict(ckpt_dict["fb"])
    paths["ba"].load_state_dict(ckpt_dict["ba"])
    paths["bb"].load_state_dict(ckpt_dict["bb"])
    paths["of"].load_state_dict(ckpt_dict["optimizer_f"])
    paths["ob"].load_state_dict(ckpt_dict["optimizer_b"])
    paths["osf"].load_state_dict(ckpt_dict["scheduler_f"])
    paths["osb"].load_state_dict(ckpt_dict["scheduler_b"])

    if device is not None:
        for opt_key in ["of", "ob"]:
            for state in paths[opt_key].state.values():
                for k, v in state.items():
                    if isinstance(v, torch.Tensor):
                        state[k] = v.to(device)

def init_parser():

    parser = argparse.ArgumentParser()
    parser.add_argument("--image_size", type=int, default=128)
    parser.add_argument("--batch_size", type=int, default=8)
    parser.add_argument("--num_epochs", type=int, default=600)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--output_path", type=str, default="/root/outputs")
    parser.add_argument("--resume_filepath", type=str, default=None)
    parser.add_argument("--seed", type=int, default=4242)
    parser.add_argument("--yaml_config", type=str, default="/root/kan-fmd/configs/LIDC.yml")

    parser.add_argument("--tr_filepath", type=str, default="")
    parser.add_argument("--tr_filepath_val", type=str, default="")
    parser.add_argument("--task_name", type=str, default="K1PUPD_SB_m2s1")

    parser.add_argument("--rsb_enabled", action="store_true", default=False,
                        help="Enable Residual Schrödinger Bridge mode")
    parser.add_argument("--lambda_bridge", type=float, default=0.1,
                        help="Bridge loss weight")
    parser.add_argument("--lambda_ddpm_anchor", type=float, default=1.0,
                        help="DDPM anchor weight for forward path")
    parser.add_argument("--lambda_alpha", type=float, default=0.01,
                        help="Alpha schedule learning weight")

    return parser
if __name__ == "__main__":

    parser = init_parser()
    args = parser.parse_args()

    set_seed(args.seed)

    train_loader, _ = get_dataloader_2(
        task="LIDC", split="train", batch_size=args.batch_size, splitratio=[0.8, 0.015, 0.185], randomsplit=True
    )    
    val_loader, _ = get_dataloader_2(
        task="LIDC", split="val", batch_size=args.batch_size, splitratio=[0.8, 0.015, 0.185], randomsplit=True
    )

    tr_dataset = TRCacheByIdDataset(
        args.tr_filepath,
        device="cpu",        # 或 "cuda"（显存够就放 cuda 更快）
        return_id=False,
        strict=True
    )
    tr_dataset_val = TRCacheByIdDataset(
        args.tr_filepath_val,
        device="cpu",        # 或 "cuda"（显存够就放 cuda 更快）
        return_id=False,
        strict=True
    )

    model = Unet(
        channels=1,
        dim_mults=(1, 2, 4),
        dim=args.image_size,
        resnet_block_groups=1,
    ).to(device)

    model = torch.compile(model, mode="reduce-overhead", fullgraph=True)

    optimizer = Adam(make_trnet_optimizer(model), lr=args.lr)

    scheduler = StepLR(
        optimizer,
        step_size=100,  # 每 100 个 epoch
        gamma=0.8       # 乘 0.8
    )

    checkpoint = None
    if args.resume_filepath is not None:
        checkpoint = torch.load(args.resume_filepath, map_location=device)
        load_model_state(model, checkpoint)
        if "optimizer_state_dict" in checkpoint:
            try:
                optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
            except ValueError:
                print("[WARN] Checkpoint optimizer does not match TRNet-only optimizer; skipping optimizer state.")
        if "scheduler_state_dict" in checkpoint:
            scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
        resume_epoch = checkpoint["epoch"]
        print(f"Resume Training on Epoch {resume_epoch}:")
    else:
        resume_epoch = 0
        print("Training from Scratch")

    from omegaconf import OmegaConf

    yaml_config = OmegaConf.load(args.yaml_config)

    def path_model_klass():
        return PathUNet(dim=args.image_size).to(device=device)
    
    paths = {"fa":path_model_klass(),
             "fb":path_model_klass(),
             "ba":path_model_klass(),
             "bb":path_model_klass(),}
    optimizer_path_f = Adam(
        list(paths['fa'].parameters()) +
        list(paths['fb'].parameters()),
        lr=1e-3
    )
    optimizer_path_b = Adam(
        list(paths['ba'].parameters()) +
        list(paths['bb'].parameters()),
        lr=1e-3
    )

    scheduler_path_f = StepLR(
        optimizer_path_f,
        step_size=50,  
        gamma=0.8       # 乘 0.8
    )
    scheduler_path_b = StepLR(
        optimizer_path_b,
        step_size=50,  
        gamma=0.8       # 乘 0.8
    )
    paths["of"] = optimizer_path_f
    paths["ob"] = optimizer_path_b
    paths["osf"] = scheduler_path_f
    paths["osb"] = scheduler_path_b

    if checkpoint is not None and checkpoint.get("paths") is not None:
        load_paths_state_rsb(paths, checkpoint["paths"], device=device)
        print("[INFO] Path models resumed from checkpoint")

    train(
        paths = paths,
        tr_dataset=tr_dataset,
        tr_dataset_val=tr_dataset_val,
        model=model,
        dataloader=train_loader,
        optimizer=optimizer,
        device=device,
        num_epochs=args.num_epochs,
        save_path=args.output_path,
        val_dataloader=val_loader,
        scheduler = scheduler,
        task_name=args.task_name,
        rsb_enabled=args.rsb_enabled,
        lambda_bridge=args.lambda_bridge,
        lambda_ddpm_anchor=args.lambda_ddpm_anchor,
        lambda_alpha=args.lambda_alpha,
    )