pretrain_PUNet.py

import os
import argparse
from datetime import datetime

import torch
from torch.optim import Adam
from torch.optim.lr_scheduler import StepLR
from tensorboardX import SummaryWriter
from tqdm import tqdm

from dataloaders import get_dataloader_2

# ProbUNet model and configuration.
from models.pretrain.ProbUNet import ProbabilisticUNet, ProbUNetConfig

# ProbUNet training and validation helpers.
from models.pretrain.probunet_train_utils import (
    train_one_epoch_prob_unet,
    validate_prob_unet,
)

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


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


def set_seed(seed):
    import random, numpy as np
    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
    os.environ["PYTHONHASHSEED"] = str(seed)
    print(f"[INFO] Global seed set to {seed}")


def train_probunet_framework(
    model: ProbabilisticUNet,
    dataloader,
    optimizer,
    device,
    num_epochs,
    save_path,
    val_dataloader,
    scheduler,
    y_mode="random1",                 # random1 / stack4 / mean_soft
    n_val_samples=16,                 # number of prior samples during validation
    thr=0.5,
    metric_mode="max_over_raters",    # max_over_raters / mean_over_raters
):
    save_path = os.path.join(save_path, "PreT_ProbUNet")
    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)

    # ProbUNet pretraining selects checkpoints by validation Dice.
    best_dice = -1.0
    best_epoch = -1

    save_epoch_list = list(range(50, num_epochs + 1, 50))

    for epoch in range(num_epochs):
        epoch_idx = epoch + 1
        print(f"Epochs:{epoch_idx}/{num_epochs} ... ")
        print("Training (ProbUNet)")

        # Train one epoch and return ELBO components.
        train_stats = train_one_epoch_prob_unet(
            model=model,
            dataloader=dataloader,
            optimizer=optimizer,
            device=device,
            y_mode=y_mode,
            bce_w=1.0,
            dice_w=1.0,
        )

        scheduler.step()

        # Validate with multiple prior samples.
        val_stats = validate_prob_unet(
            model=model,
            dataloader=val_dataloader,
            device=device,
            n_samples=n_val_samples,
            thr=thr,
            metric_mode=metric_mode,
        )

        print(
            f"train: epoch:{epoch_idx:04d} "
            f"loss:{train_stats['loss']:.4f} seg:{train_stats['seg']:.4f} kl:{train_stats['kl']:.4f} | "
            f"val: dice:{val_stats['dice']:.4f} iou:{val_stats['iou']:.4f}"
        )

        # TensorBoardX
        writer.add_scalar("train/loss", train_stats["loss"], epoch_idx)
        writer.add_scalar("train/seg",  train_stats["seg"],  epoch_idx)
        writer.add_scalar("train/kl",   train_stats["kl"],   epoch_idx)

        writer.add_scalars("val/metrics", {"Dice": val_stats["dice"], "IoU": val_stats["iou"]}, epoch_idx)

        # Save best checkpoint by validation Dice.
        if val_stats["dice"] > best_dice:
            best_dice = val_stats["dice"]
            best_epoch = epoch_idx
            print(f"Saving best model: epoch:{epoch_idx:04d}-Dice:{best_dice:.4f}")

            checkpoint = {
                "epoch": epoch_idx,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "scheduler_state_dict": scheduler.state_dict(),
                "best_dice": best_dice,
                "cfg": model.cfg.__dict__,   # saved for reproducibility
            }
            os.makedirs(save_path, exist_ok=True)
            torch.save(checkpoint, os.path.join(save_path, "best.pth"))

        # Periodic checkpointing.
        if epoch_idx in save_epoch_list:
            checkpoint = {
                "epoch": epoch_idx,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "scheduler_state_dict": scheduler.state_dict(),
                "cfg": model.cfg.__dict__,
            }
            os.makedirs(save_path, exist_ok=True)
            torch.save(checkpoint, os.path.join(save_path, f"epoch_{epoch_idx:04d}.pth"))

    print(f"Best model: epoch:{best_epoch:04d}-Dice:{best_dice:.4f}")
    writer.close()


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=1600)
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)

# ProbUNet-specific options.
parser.add_argument("--z_dim", type=int, default=6)
parser.add_argument("--beta", type=float, default=1.0)
parser.add_argument("--y_mode", type=str, default="random1", choices=["random1", "stack4", "mean_soft"])
parser.add_argument("--val_samples", type=int, default=16)
parser.add_argument("--metric_mode", type=str, default="max_over_raters", choices=["max_over_raters", "mean_over_raters"])


if __name__ == "__main__":
    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.1, 0.1], randomsplit=True
    )
    val_loader, _ = get_dataloader_2(
        task="LIDC", split="val", batch_size=args.batch_size, splitratio=[0.8, 0.1, 0.1], randomsplit=True
    )

    # ProbUNet model.
    model_cfg = ProbUNetConfig(
        in_channels=1,
        out_channels=1,
        z_dim=args.z_dim,
        beta=args.beta,

        unet_base_channels=32,
        unet_depth=4,
        norm="gn",
        dropout=0.0,
        act="relu",
        up_mode="bilinear",

        latent_base_channels=32,
        latent_depth=4,
        fcomb_hidden=64,
        fcomb_layers=2,
    )
    model = ProbabilisticUNet(model_cfg).to(device)

    # Optional compile path; disabled by default for easier debugging.
    # model = torch.compile(model, mode="reduce-overhead", fullgraph=True)

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

    # resume
    if args.resume_filepath is not None:
        checkpoint = torch.load(args.resume_filepath, map_location="cpu")
        model.load_state_dict(checkpoint["model_state_dict"])
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        scheduler.load_state_dict(checkpoint.get("scheduler_state_dict", scheduler.state_dict()))
        resume_epoch = checkpoint.get("epoch", 0)
        print(f"Resume Training on Epoch {resume_epoch}:")
    else:
        resume_epoch = 0
        print("Training from Scratch")

    train_probunet_framework(
        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,
        y_mode=args.y_mode,
        n_val_samples=args.val_samples,
        metric_mode=args.metric_mode,
    )