train_stage2_trnet.py

import os
import argparse
import random
import numpy as np

import torch
from torch.optim import Adam
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader
from tqdm import tqdm
from tensorboardX import SummaryWriter
from omegaconf import OmegaConf
from datetime import datetime

from models.kan_resdiff_sb import Unet
from models.path_sb.diffusion import PathUNet
from models.path_sb.rsb import bridge_loss
from models.pretrain.TRDataset import TRCacheByIdDataset
from dataloaders import *
from metrics import *
from validate import validate
from checkpoint_utils import load_model_state

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# ---------------------------------------------------------------------------
# utility functions
# ---------------------------------------------------------------------------

def set_seed(seed):
    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)

def make_timestamp_dir(save_path):
    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    return os.path.join(save_path, timestamp)

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 make_ddpm_schedule(T=1000, beta_start=1e-4, beta_end=2e-2, device="cuda", num_timesteps=None):
    if num_timesteps is not None:
        T = num_timesteps
    betas = torch.linspace(beta_start, beta_end, T, device=device)
    alphas = 1.0 - betas
    alphas_bar = torch.cumprod(alphas, dim=0)
    sqrt_ab = torch.sqrt(alphas_bar)
    sqrt_1mab = torch.sqrt(1.0 - alphas_bar)
    return betas, alphas, alphas_bar, sqrt_ab, sqrt_1mab

def solve_a_b(x0, eps, x_pred, eps_stability=1e-8):
    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 set_stage2_model_trainability(model, train_trnet):
    """Stage 2 keeps the diffusion backbone frozen; only TRNet may train."""
    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 the TRNet-only parameter list 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 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):
    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 save_topk_models(best_models, top_k, metrics, checkpoint, save_path):
    filename = (
        f"Best@Epoch:{metrics['epoch']:04d}_"
        f"MDM:{metrics['MDM']:.4f}_GED{metrics['GED']:.4f}_"
        f"IoU:{metrics['IoU']:.4f}.pth"
    )
    ckpt_path = os.path.join(save_path, filename)

    if len(best_models) < top_k:
        torch.save(checkpoint, ckpt_path)
        best_models.append({**metrics, "path": ckpt_path})
    else:
        worst = max(best_models, key=lambda x: x["GED"])
        if metrics["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({**metrics, "path": ckpt_path})
    best_models.sort(key=lambda x: x["GED"])

# ---------------------------------------------------------------------------
# stage-2 training loop
# ---------------------------------------------------------------------------

def train_stage2(
    paths,
    tr_dataset,
    tr_dataset_val,
    model,
    dataloader,
    optimizer,
    device,
    num_epochs,
    save_path,
    val_dataloader,
    scheduler,
    task_name,
    lambda_bridge=0.1,
    lambda_ddpm_anchor=1.0,
    lambda_alpha=0.01,
):
    TOP_K = 3
    best_models = []
    num_timesteps = 1000
    _, _, _, sqrt_ab, sqrt_1mab = make_ddpm_schedule(T=num_timesteps, device=device)

    save_dir = os.path.join(save_path, task_name)
    save_dir = make_timestamp_dir(save_dir)
    os.makedirs(save_dir, exist_ok=True)

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

    resume_loaded_paths = False
    if args.resume_filepath is not None:
        ckpt_resume = torch.load(args.resume_filepath, map_location=device)
        if "paths" in ckpt_resume and ckpt_resume["paths"] is not None:
            load_paths_state(paths, ckpt_resume["paths"], device=device)
            resume_loaded_paths = True
            print(f"Stage 2 paths resumed from {args.resume_filepath}")

    if not resume_loaded_paths:
        stage1_path = args.stage1_filepath or os.path.join(save_path, "paths_stage1.pth")
        if not os.path.exists(stage1_path):
            raise FileNotFoundError(
                f"Stage 1 paths not found: {stage1_path}\n"
                f"Run train_stage1_path.py first."
            )
        ckpt_s1 = torch.load(stage1_path, map_location=device)
        load_paths_state(paths, ckpt_s1["paths"], device=device)
        print(f"Stage 1 paths loaded from {stage1_path}")

    for k in ["fa", "fb", "ba", "bb"]:
        paths[k].train()
        for p in paths[k].parameters():
            p.requires_grad = True

    for epoch in range(num_epochs):
        is_odd = (epoch % 2 == 1)

        if is_odd:
            model._stage(1)
            set_stage2_model_trainability(model, train_trnet=True)
            print(f"{task_name}; Epochs:{epoch + 1}/{num_epochs} [2n+1: TRNet + backward path]")
        else:
            set_stage2_model_trainability(model, train_trnet=False)
            model.eval()
            print(f"{task_name}; Epochs:{epoch + 1}/{num_epochs} [2n+2: forward path]")

        total_loss = 0.0
        total_fwd = 0.0
        total_bwd = 0.0
        total_bridge = 0.0

        pbar = tqdm(dataloader, ncols=100, leave=False)
        for images, masks, _, case_ids in pbar:
            if is_odd:
                optimizer.zero_grad()
            paths["of"].zero_grad()
            paths["ob"].zero_grad()

            images = images.to(device)
            batch_size = images.shape[0]
            m1, m2, m3, m4 = [m.to(device) for m in masks]
            all_masks = torch.cat([m1, m2, m3, m4], dim=0)
            all_images = images.repeat(4, 1, 1, 1)

            t_int = torch.randint(0, num_timesteps, (batch_size,), device=device, dtype=torch.long)
            t_batch = t_int.repeat(4)
            t_norm = t_batch.float() / (num_timesteps - 1)
            eps = torch.randn_like(all_masks)

            a_f = paths['fa'](eps, t_norm, all_images).view(-1)
            b_f = paths['fb'](eps, t_norm, all_images).view(-1)
            a_ref = sqrt_ab[t_batch]
            b_ref = sqrt_1mab[t_batch]

            if is_odd:
                # ---------- 2n+1: TRNet training + backward path update ----------
                a_f_det = a_f.detach().view(-1, 1, 1, 1)
                b_f_det = b_f.detach().view(-1, 1, 1, 1)

                x_t_tr = a_f_det * all_masks + b_f_det * eps

                TR_list = [tr_dataset.get_by_id(cid) for cid in case_ids]
                TR = torch.stack(TR_list, dim=0).to(device, non_blocking=True)
                TR_rep = TR.repeat(4, 1, 1, 1)

                # Step A: train TRNet on the learned-path state.
                optimizer.zero_grad()
                eps_pred = model(x_t_tr, t_norm, all_images, TR_rep, x_tr=x_t_tr).float()

                x0_hat = (x_t_tr - b_f_det * eps_pred) / (a_f_det + 1e-8)
                x0_hat = x0_hat.clamp(0, 1)

                L_eps = (eps_pred - eps).pow(2).mean()
                dice_val = dice_batch(x0_hat, all_masks)
                L_dice = (1 - dice_val).mean()
                L_data = L_eps + 0.0001 * L_dice

                L_alpha = torch.tensor(0.0, device=device)
                if model.alpha is not None:
                    L_alpha = lambda_alpha * ((model.alpha - (1.0 - t_norm)) ** 2).mean()

                L_tr = L_data + L_alpha
                L_tr.backward()
                optimizer.step()

                # Step B: solve the coefficients implied by the updated TRNet prediction.
                with torch.no_grad():
                    eps_pred_new = model(x_t_tr, t_norm, all_images, TR_rep, x_tr=x_t_tr).float()
                a_model, b_model = solve_a_b(x0=all_masks, eps=eps_pred_new, x_pred=x_t_tr)
                a_b = paths['ba'](eps_pred_new, t_norm, all_images).view(-1)
                b_b = paths['bb'](eps_pred_new, t_norm, all_images).view(-1)

                # Step C: train the backward path to match those implied coefficients.
                paths["ob"].zero_grad()
                L_bwd_path = ((a_b - a_model) ** 2).mean() + ((b_b - b_model) ** 2).mean()
                L_bwd_path.backward()
                paths["ob"].step()

                L_total = L_tr + L_bwd_path.detach()
                total_bwd += L_data.item()
                total_bridge += L_bwd_path.item()

                postfix = {
                    "loss": f"{L_total.item():.4f}",
                    "TR_data": f"{L_data.item():.4f}",
                    "bwd_path": f"{L_bwd_path.item():.4f}",
                }
            else:
                # ---------- 2n+2: forward path update ----------
                with torch.no_grad():
                    a_f_det = a_f.view(-1, 1, 1, 1)
                    b_f_det = b_f.view(-1, 1, 1, 1)
                    x_t_bridge = a_f_det * all_masks + b_f_det * eps
                    eps_pred_bridge = model(x_t_bridge, t_norm, all_images, None).float()
                    a_b_bridge = paths['ba'](eps_pred_bridge, t_norm, all_images).view(-1)
                    b_b_bridge = paths['bb'](eps_pred_bridge, t_norm, all_images).view(-1)

                L_bridge = lambda_bridge * bridge_loss(a_f, b_f, a_b_bridge.detach(), b_b_bridge.detach())
                L_anchor = lambda_ddpm_anchor * (
                    ((a_f - a_ref) ** 2).mean() + ((b_f - b_ref) ** 2).mean()
                )

                L_total = L_bridge + L_anchor
                L_total.backward()
                paths["of"].step()

                total_fwd += L_anchor.item()
                total_bridge += L_bridge.item()

                postfix = {
                    "loss": f"{L_total.item():.4f}",
                    "fwd_anchor": f"{L_anchor.item():.4f}",
                    "bridge": f"{L_bridge.item():.4f}",
                }

            total_loss += L_total.item()
            pbar.set_postfix(postfix)

        if is_odd:
            scheduler.step()
            paths["osb"].step()
        else:
            paths["osf"].step()
        epoch_idx = epoch + 1
        avg_loss = total_loss / len(dataloader)
        avg_bridge = total_bridge / len(dataloader)
        print(f"[Epoch {epoch_idx}] Loss={avg_loss:.6f} Bridge={avg_bridge:.6f}")

        mdm, ged, iou = validate(
            tr_dataset_val, model, val_dataloader, device,
            sampling_times=24, num_samples=4, print_full=False, tr_enabled=True,
            paths=paths, sampling_mode="rsb",
        )
        print(f"Epoch:{epoch_idx:04d} MDM:{mdm:.4f} GED:{ged:.4f} IoU:{iou:.4f}")

        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(paths),
            "GED": ged,
        }
        save_topk_models(best_models, TOP_K, metrics, checkpoint, save_dir)

        if epoch_idx % 50 == 0:
            ckpt_name = f"Epoch:{epoch_idx:04d}-MDM:{mdm:.4f}-GED:{ged:.4f}-IoU:{iou:.4f}.pth"
            torch.save(checkpoint, os.path.join(save_dir, ckpt_name))

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

# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    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("--stage1_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("--lambda_bridge", type=float, default=0.1)
    parser.add_argument("--lambda_ddpm_anchor", type=float, default=1.0)
    parser.add_argument("--lambda_alpha", type=float, default=0.01)
    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", return_id=False, strict=True)
    tr_dataset_val = TRCacheByIdDataset(args.tr_filepath_val, device="cpu", 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, gamma=0.8)

    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"])
        print(f"Resume Training on Epoch {checkpoint['epoch']}")

    OmegaConf.load(args.yaml_config)

    def make_path_unet():
        return PathUNet(dim=args.image_size).to(device=device)

    paths = {"fa": make_path_unet(), "fb": make_path_unet(),
             "ba": make_path_unet(), "bb": make_path_unet()}

    opt_f = Adam(list(paths['fa'].parameters()) + list(paths['fb'].parameters()), lr=1e-3)
    opt_b = Adam(list(paths['ba'].parameters()) + list(paths['bb'].parameters()), lr=1e-3)

    paths["of"] = opt_f
    paths["ob"] = opt_b
    paths["osf"] = StepLR(opt_f, step_size=50, gamma=0.8)
    paths["osb"] = StepLR(opt_b, step_size=50, gamma=0.8)

    train_stage2(
        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,
        lambda_bridge=args.lambda_bridge,
        lambda_ddpm_anchor=args.lambda_ddpm_anchor,
        lambda_alpha=args.lambda_alpha,
    )