train_stage1_path.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 dataloaders import *
from metrics import *
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 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

@torch.no_grad()
def sample_model_path(model, mask_gt, img_cond, t_int, sqrt_ab, sqrt_1mab):
    N = mask_gt.shape[0]
    num_timesteps = sqrt_ab.shape[0]

    if isinstance(t_int, int):
        t_batch = torch.full((N,), t_int, device=mask_gt.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() / (num_timesteps - 1)

    eps_pred = model(x_t, t_norm, img_cond, None).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 {
        "eps": eps,
        "eps_pred": eps_pred.detach(),
        "x0_hat": x0_hat,
        "a_model": a_model,
        "b_model": b_model,
    }

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(),
    }

# ---------------------------------------------------------------------------
# stage-1 path initialisation
# ---------------------------------------------------------------------------

def train_stage1(paths, model, dataloader, device, num_epochs, save_path, task_name):
    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)

    model._stage(0)
    model.eval()
    best_loss = float("inf")

    for epoch in range(num_epochs):
        print(f"{task_name}; Epochs:{epoch + 1}/{num_epochs} [Stage 1: Path Init]")

        total_fwd = 0.0
        total_bwd = 0.0

        pbar = tqdm(dataloader, ncols=100, leave=False)
        for images, masks, _, _ in pbar:
            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)

            result = sample_model_path(
                model=model, mask_gt=all_masks, img_cond=all_images,
                t_int=t_batch, sqrt_ab=sqrt_ab, sqrt_1mab=sqrt_1mab,
            )

            a_ref = sqrt_ab[t_batch]
            b_ref = sqrt_1mab[t_batch]
            a_model = result['a_model'].detach()
            b_model = result['b_model'].detach()

            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_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)

            loss_fwd = ((a_f - a_ref) ** 2).mean() + ((b_f - b_ref) ** 2).mean()
            loss_bwd = ((a_b - a_model) ** 2).mean() + ((b_b - b_model) ** 2).mean()

            paths["of"].zero_grad()
            loss_fwd.backward()
            paths["of"].step()

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

            total_fwd += loss_fwd.item()
            total_bwd += loss_bwd.item()
            pbar.set_postfix({"fwd": f"{loss_fwd.item():.4f}", "bwd": f"{loss_bwd.item():.4f}"})

        epoch_idx = epoch + 1
        paths["osf"].step()
        paths["osb"].step()
        avg_fwd = total_fwd / len(dataloader)
        avg_bwd = total_bwd / len(dataloader)
        combined = avg_fwd + avg_bwd
        print(f"[Epoch {epoch_idx}] Fwd={avg_fwd:.6f}  Bwd={avg_bwd:.6f}")

        if combined < best_loss:
            best_loss = combined
            torch.save(
                {"epoch": epoch_idx, "model": model.state_dict(), "paths": get_paths_state(paths)},
                os.path.join(save_dir, "paths_stage1.pth"),
            )
            print(f"Best saved (combined={combined:.6f})")

    writer.close()
    print("Stage 1 complete.")

# ---------------------------------------------------------------------------
# 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("--seed", type=int, default=4242)
    parser.add_argument("--yaml_config", type=str, default="/root/kan-fmd/configs/LIDC.yml")
    parser.add_argument("--task_name", type=str, default="K1PUPD_SB_stage1")
    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,
    )

    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(model.parameters(), lr=args.lr)

    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:
            optimizer.load_state_dict(checkpoint["optimizer_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_stage1(
        paths=paths, model=model, dataloader=train_loader,
        device=device, num_epochs=args.num_epochs,
        save_path=args.output_path, task_name=args.task_name,
    )