train_stage0_backbone.py

import argparse
import os
import random
from datetime import datetime

import numpy as np
import torch
from omegaconf import OmegaConf
from tensorboardX import SummaryWriter
from torch.optim import Adam
from torch.optim.lr_scheduler import StepLR
from tqdm import tqdm

from checkpoint_utils import load_model_state
from dataloaders import *
from metrics import *
from models.kan_resdiff_sb import Unet
from validate import validate


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


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


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)
    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 loss_fn_ddpm(model, mask_gt, t_int, img_cond, 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)

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

    x0_hat = (x_t - b * eps_pred) / (a + 1e-8)
    y_bin = (x0_hat.clamp(0, 1) >= 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_stage0(model, dataloader, optimizer, device, num_epochs, save_path, val_dataloader, scheduler, task_name):
    top_k = 3
    best_models = []
    T = 1000
    _, _, _, sqrt_ab, sqrt_1mab = make_ddpm_schedule(T=T, device=device)

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

    model._stage(0)

    for epoch in range(num_epochs):
        model.train()
        model._stage(0)
        total_loss = 0.0
        total_dice = 0.0
        total_iou = 0.0

        pbar = tqdm(dataloader, ncols=100, leave=False)
        for images, masks, _, _ in pbar:
            optimizer.zero_grad()
            images = images.to(device)
            batch_size = images.shape[0]
            all_masks = torch.cat([m.to(device) for m in masks], dim=0)
            all_images = images.repeat(4, 1, 1, 1)

            t_int = torch.randint(0, T, (batch_size,), device=device, dtype=torch.long).repeat(4)
            loss_all, dice_all, iou_all = loss_fn_ddpm(
                model=model,
                mask_gt=all_masks,
                t_int=t_int,
                img_cond=all_images,
                sqrt_ab=sqrt_ab,
                sqrt_1mab=sqrt_1mab,
            )

            loss = loss_all.mean()
            dice = dice_all.mean()
            iou = iou_all.mean()
            back_loss = loss + 0.0001 * (1 - dice)
            back_loss.backward()
            optimizer.step()

            total_loss += loss.item()
            total_dice += dice.item()
            total_iou += iou.item()
            pbar.set_postfix({"loss": f"{loss.item():.4f}", "dice": f"{dice.item():.4f}", "iou": f"{iou.item():.4f}"})

        scheduler.step()
        epoch_idx = epoch + 1
        avg_loss = total_loss / len(dataloader)
        avg_dice = total_dice / len(dataloader)
        avg_iou = total_iou / len(dataloader)
        print(f"Epoch:{epoch_idx}, Loss={avg_loss:.6f}, Dice={avg_dice:.6f}, IoU={avg_iou:.6f}")

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

        writer.add_scalar("train/Loss", avg_loss, epoch_idx)
        writer.add_scalars("train/metrics", {"Dice": avg_dice, "IoU": avg_iou}, epoch_idx)
        writer.add_scalars("val/metrics", {"MDM": mdm, "GED": ged, "IoU": iou}, epoch_idx)

        checkpoint = {
            "epoch": epoch_idx,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "scheduler_state_dict": scheduler.state_dict(),
            "loss": avg_loss,
            "GED": ged,
        }
        metrics = {"epoch": epoch_idx, "MDM": mdm, "GED": ged, "IoU": iou}
        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}")


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("--task_name", type=str, default="K1PUPD_SB_stage0")
    return parser


if __name__ == "__main__":
    args = init_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,
    )

    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)
    model._stage(0)

    optimizer = Adam(model.parameters(), lr=args.lr)
    scheduler = StepLR(optimizer, step_size=100, gamma=0.8)

    resume_epoch = 0
    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"])
        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}:")

    OmegaConf.load(args.yaml_config)

    train_stage0(
        model=model,
        dataloader=train_loader,
        optimizer=optimizer,
        device=device,
        num_epochs=args.num_epochs - resume_epoch,
        save_path=args.output_path,
        val_dataloader=val_loader,
        scheduler=scheduler,
        task_name=args.task_name,
    )