sacflow/engine/train_loop.py

from __future__ import annotations
from pathlib import Path
import time
import copy
import re
import torch
import torch.nn.functional as F
from torch.nn.parallel import DistributedDataParallel as DDP
from tqdm import tqdm
from monai.inferers import sliding_window_inference
from sacflow.data.loader import build_loader
from sacflow.models.unet3d import build_model, freeze_except_adapters
from sacflow.models.velocity_field import VelocityField3D
from sacflow.methods.sacflow_step import sacflow_forward_step, ce_loss_masked, dice_loss_masked
from sacflow.methods.source_memory import load_source_memory, class_moments
from sacflow.utils.metrics import torch_soft_dice_loss, entropy_loss, confidence_and_margin, dice_per_class, hd95_per_class
from sacflow.utils.misc import ensure_dir, count_trainable, move_to_device, unwrap_model
from sacflow.utils.distributed import is_main_process, get_world_size, get_rank, reduce_mean, barrier, is_dist_avail_and_initialized
import torch.distributed as dist
from sacflow.utils.wandb_utils import wandb_log


def build_optimizer(params, cfg):
    ocfg = cfg["optim"]
    params = [p for p in params if p.requires_grad]
    if ocfg.get("optimizer", "adamw").lower() == "sgd":
        return torch.optim.SGD(params, lr=float(ocfg["lr"]), momentum=0.9, weight_decay=float(ocfg.get("weight_decay", 0)))
    return torch.optim.AdamW(params, lr=float(ocfg["lr"]), weight_decay=float(ocfg.get("weight_decay", 0)), betas=tuple(ocfg.get("betas", [0.9, 0.999])))


def update_ema(teacher, student, decay):
    with torch.no_grad():
        for pt, ps in zip(teacher.parameters(), student.parameters()):
            pt.data.mul_(decay).add_(ps.data, alpha=1-decay)


def load_checkpoint_into(model, path, strict=False):
    ckpt = torch.load(path, map_location="cpu")
    state = ckpt.get("model", ckpt)
    missing, unexpected = model.load_state_dict(state, strict=strict)
    return missing, unexpected


def save_checkpoint(path, model, optimizer, epoch, step, best_metric=None, velocity_field=None, cfg=None, teacher=None, include_optimizer=True):
    """Save a checkpoint on rank 0 only.

    Disk policy:
    - best.pt is intended for evaluation/inference and is saved without optimizer by default.
    - last.pt is intended for resume and includes optimizer.
    This avoids filling the disk with epoch_N.pt checkpoints.
    """
    if not is_main_process():
        return
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    obj = {
        "model": unwrap_model(model).state_dict(),
        "epoch": epoch,
        "step": step,
        "best_metric": best_metric,
        "cfg": cfg,
    }
    if include_optimizer and optimizer is not None:
        obj["optimizer"] = optimizer.state_dict()
    if velocity_field is not None:
        obj["velocity_field"] = unwrap_model(velocity_field).state_dict()
    if teacher is not None:
        obj["teacher"] = unwrap_model(teacher).state_dict()
    torch.save(obj, path)


def _epoch_number(path: Path) -> int:
    m = re.search(r"epoch_(\d+)\.pt$", path.name)
    return int(m.group(1)) if m else -1


def resolve_resume_checkpoint(ckpt_dir: Path, resume_value):
    """Return a usable resume checkpoint path.

    resume_value can be:
    - None/False: do not resume
    - "auto"/True: prefer last.pt, then newest epoch_*.pt, then best.pt
    - explicit checkpoint path
    Corrupted/incomplete checkpoints are skipped.
    """
    if not resume_value:
        return None
    if str(resume_value).lower() not in ("auto", "true", "1", "yes"):
        return Path(resume_value)
    candidates = []
    last = ckpt_dir / "last.pt"
    if last.exists():
        candidates.append(last)
    candidates.extend(sorted(ckpt_dir.glob("epoch_*.pt"), key=_epoch_number, reverse=True))
    best = ckpt_dir / "best.pt"
    if best.exists():
        candidates.append(best)
    for c in candidates:
        try:
            torch.load(c, map_location="cpu")
            return c
        except Exception as e:
            if is_main_process():
                print(f"Skipping unusable checkpoint {c}: {e}")
    return None


def load_training_checkpoint(path, model, optimizer=None, velocity_field=None, teacher=None):
    ckpt = torch.load(path, map_location="cpu")
    missing, unexpected = unwrap_model(model).load_state_dict(ckpt.get("model", ckpt), strict=False)
    if is_main_process():
        print(f"Loaded resume model from {path} missing={len(missing)} unexpected={len(unexpected)}")
    if velocity_field is not None and "velocity_field" in ckpt:
        unwrap_model(velocity_field).load_state_dict(ckpt["velocity_field"], strict=False)
    if teacher is not None and "teacher" in ckpt:
        unwrap_model(teacher).load_state_dict(ckpt["teacher"], strict=False)
    elif teacher is not None:
        unwrap_model(teacher).load_state_dict(unwrap_model(model).state_dict(), strict=False)
    if optimizer is not None and ckpt.get("optimizer") is not None:
        optimizer.load_state_dict(ckpt["optimizer"])
    start_epoch = int(ckpt.get("epoch", 0))
    global_step = int(ckpt.get("step", 0))
    best = float(ckpt.get("best_metric", -1e9) if ckpt.get("best_metric", None) is not None else -1e9)
    return start_epoch, global_step, best


def supervised_step(model, batch, cfg):
    x = batch["image"]
    y = batch["label"]
    logits = model(x)
    ce = F.cross_entropy(logits, y.long())
    dice = torch_soft_dice_loss(logits, y, cfg["data"]["num_classes"])
    loss = cfg["train"].get("loss", {}).get("ce", 1.0)*ce + cfg["train"].get("loss", {}).get("dice", 1.0)*dice
    return loss, {"loss_total": loss.detach(), "loss_ce": ce.detach(), "loss_dice": dice.detach()}



def proto_align_step(model, teacher, batch, memory, cfg):
    x = batch["image"]
    logits, feats = model(x, return_features=True)
    feat = feats["prelogit"]
    with torch.no_grad():
        tlogits = teacher(x)
        tprobs = torch.softmax(tlogits, dim=1)
        conf, margin, pseudo = confidence_and_margin(tprobs)
        mask = conf > float(cfg["train"].get("pseudo_conf_threshold", 0.75))
        probs_f = tprobs
        if probs_f.shape[-3:] != feat.shape[-3:]:
            probs_f = F.interpolate(probs_f, size=feat.shape[-3:], mode="trilinear", align_corners=False)
    ce = ce_loss_masked(logits, pseudo, mask)
    dice = dice_loss_masked(logits, pseudo, mask, cfg["data"]["num_classes"])
    proto_loss = torch.tensor(0.0, device=x.device)
    if memory is not None and "feature_mu" in memory:
        mu = memory["feature_mu"].to(feat.device, feat.dtype)  # [C,d]
        # expected source prototype at each voxel based on teacher probabilities
        proto = torch.einsum("bchwz,cf->bfhwz", probs_f.detach(), mu)
        proto_loss = ((feat - proto).pow(2) * probs_f.max(1, keepdim=True).values.detach()).mean()
    ent = entropy_loss(logits)
    loss_cfg = cfg["train"].get("loss", {})
    loss = float(loss_cfg.get("ce", 1.0))*ce + float(loss_cfg.get("dice", 1.0))*dice + float(loss_cfg.get("prototype", 0.1))*proto_loss + float(loss_cfg.get("entropy", 0.01))*ent
    return loss, {"loss_total": loss.detach(), "loss_pseudo_ce": ce.detach(), "loss_pseudo_dice": dice.detach(), "loss_proto_align": proto_loss.detach(), "loss_entropy": ent.detach(), "pseudo_conf_mean": conf.mean().detach(), "pseudo_accept_rate": mask.float().mean().detach()}


def pseudo_step(model, teacher, batch, cfg):
    x = batch["image"]
    with torch.no_grad():
        tlogits = teacher(x)
        tprobs = torch.softmax(tlogits, dim=1)
        conf, margin, pseudo = confidence_and_margin(tprobs)
        mask = conf > float(cfg["train"].get("pseudo_conf_threshold", 0.75))
    logits = model(x)
    ce = ce_loss_masked(logits, pseudo, mask)
    dice = dice_loss_masked(logits, pseudo, mask, cfg["data"]["num_classes"])
    ent = entropy_loss(logits)
    loss_cfg = cfg["train"].get("loss", {})
    loss = float(loss_cfg.get("ce", 1.0))*ce + float(loss_cfg.get("dice", 1.0))*dice + float(loss_cfg.get("entropy", 0.01))*ent
    return loss, {"loss_total": loss.detach(), "loss_pseudo_ce": ce.detach(), "loss_pseudo_dice": dice.detach(), "loss_entropy": ent.detach(), "pseudo_conf_mean": conf.mean().detach(), "pseudo_accept_rate": mask.float().mean().detach()}


@torch.no_grad()
def evaluate(model, loader, cfg, device, max_batches=None):
    """Evaluate segmentation metrics.

    In DDP this function is called on *all* ranks with a no-padding sharded
    validation loader. It then all-reduces metric sums/counts so rank 0 gets
    exact full-validation metrics without other ranks idling at a barrier.
    """
    model.eval()
    all_metrics = []
    num_classes = cfg["data"]["num_classes"]
    roi_size = tuple(cfg.get("eval", {}).get("roi_size", cfg["data"].get("patch_size", [96,96,96])))
    sw_batch_size = int(cfg.get("eval", {}).get("sw_batch_size", 1))
    overlap = float(cfg.get("eval", {}).get("overlap", 0.5))
    iterator = enumerate(loader)
    if is_main_process():
        iterator = tqdm(iterator, total=len(loader), desc="eval", leave=False)
    for i, batch in iterator:
        if max_batches is not None and i >= max_batches:
            break
        if "label" not in batch:
            continue
        x = batch["image"].to(device, non_blocking=True)
        y = batch["label"].numpy()
        if cfg.get("eval", {}).get("sliding_window", True):
            logits = sliding_window_inference(x, roi_size=roi_size, sw_batch_size=sw_batch_size, predictor=model, overlap=overlap)
        else:
            logits = model(x)
        pred = logits.argmax(1).cpu().numpy()
        for b in range(pred.shape[0]):
            m = {}
            m.update(dice_per_class(pred[b], y[b], num_classes))
            spacing = tuple(batch.get("spacing", torch.ones(1,3))[b].cpu().numpy().tolist()) if "spacing" in batch else (1,1,1)
            m.update(hd95_per_class(pred[b], y[b], num_classes, spacing=spacing))
            all_metrics.append(m)

    metric_keys = [f"dice_c{c}" for c in range(1, num_classes)] + ["dice_mean"] + [f"hd95_c{c}" for c in range(1, num_classes)] + ["hd95_mean"]
    sums = torch.zeros(len(metric_keys), device=device, dtype=torch.float64)
    counts = torch.zeros(len(metric_keys), device=device, dtype=torch.float64)
    for m in all_metrics:
        for j, k in enumerate(metric_keys):
            v = m.get(k, float("nan"))
            if v == v:  # not NaN
                sums[j] += float(v)
                counts[j] += 1.0
    if is_dist_avail_and_initialized():
        dist.all_reduce(sums, op=dist.ReduceOp.SUM)
        dist.all_reduce(counts, op=dist.ReduceOp.SUM)
    out = {}
    for j, k in enumerate(metric_keys):
        if counts[j].item() > 0:
            out[f"val/{k}"] = float((sums[j] / counts[j]).item())
    if not out:
        out["val/dice_mean"] = float("nan")
    return out


def run_training(cfg, device, wandb_run=None):
    mode = cfg["train"]["mode"]
    out_dir = ensure_dir(cfg["output_dir"])
    ckpt_dir = ensure_dir(out_dir / "checkpoints")
    require_label = mode in ("source_train", "oracle_train")
    split = "source_train" if mode == "source_train" else ("target_train" if mode in ("oracle_train", "self_train", "peft", "sacflow_fm", "proto_align") else "target_train")
    train_loader = build_loader(cfg, split=split, training=True, require_label=require_label)
    val_split = "source_val" if mode == "source_train" else "target_val"
    try:
        val_loader = build_loader(cfg, split=val_split, training=False, require_label=True, distributed=(get_world_size() > 1))
    except Exception:
        val_loader = None
    model = build_model(cfg).to(device)
    if cfg["train"].get("source_checkpoint"):
        missing, unexpected = load_checkpoint_into(model, cfg["train"]["source_checkpoint"], strict=False)
        if is_main_process():
            print("Loaded source checkpoint", cfg["train"]["source_checkpoint"], "missing", len(missing), "unexpected", len(unexpected))
    if mode in ("peft", "sacflow_fm", "proto_align") and cfg.get("model", {}).get("adapter", {}).get("enabled", False):
        freeze_except_adapters(model, train_norm_affine=True)
    teacher = copy.deepcopy(model).to(device)
    for p in teacher.parameters():
        p.requires_grad = False
    velocity_field = None
    memory = None
    if mode in ("sacflow_fm", "proto_align"):
        if cfg["train"].get("memory_path"):
            memory = load_source_memory(cfg["train"]["memory_path"], device=device)
        feat_ch = model.prelogit_channels
        vcfg = cfg.get("sacflow", {}).get("velocity", {})
        if cfg.get("sacflow", {}).get("use_velocity_field", True):
            velocity_field = VelocityField3D(
                residual_channels=feat_ch,
                num_classes=cfg["data"]["num_classes"],
                hidden_ratio=float(vcfg.get("hidden_ratio", 0.25)),
                depth=int(vcfg.get("depth", 2)),
                tau_embedding_dim=int(vcfg.get("tau_embedding_dim", 32)),
                organ_embedding_dim=int(vcfg.get("organ_embedding_dim", 16)),
                include_teacher_probs=bool(vcfg.get("include_teacher_probs", True)),
                include_confidence=bool(vcfg.get("include_confidence", True)),
                include_boundary=bool(vcfg.get("include_boundary", True)),
                use_depthwise=bool(vcfg.get("use_depthwise", True)),
                use_group_norm=bool(vcfg.get("use_group_norm", True)),
                use_film=bool(vcfg.get("use_film", True)),
            ).to(device)
    params = list(model.parameters()) + ([] if velocity_field is None else list(velocity_field.parameters()))
    optimizer = build_optimizer(params, cfg)
    if get_world_size() > 1:
        # SACFlow uses a custom feature-path loss in addition to the ordinary forward pass.
        # find_unused_parameters=True is safer for this mode because some classifier outputs
        # from the first forward are not directly used in the loss, while PEFT adapters are
        # used again for path-state classification.
        find_unused = bool(cfg.get("distributed", {}).get("find_unused_parameters", False)) or mode == "sacflow_fm"
        model = DDP(model, device_ids=[device.index] if device.type == "cuda" else None, find_unused_parameters=find_unused)
        if velocity_field is not None:
            velocity_field = DDP(velocity_field, device_ids=[device.index] if device.type == "cuda" else None, find_unused_parameters=True)
    trainable, total = count_trainable(unwrap_model(model))
    if velocity_field is not None:
        vt, vtotal = count_trainable(unwrap_model(velocity_field))
        trainable += vt
        total += vtotal
    if is_main_process():
        print(f"Mode={mode} trainable={trainable:,} total={total:,} ({100*trainable/max(1,total):.2f}%)")
    scaler = torch.cuda.amp.GradScaler(enabled=bool(cfg.get("amp", True)) and device.type == "cuda")
    best = -1e9
    global_step = 0
    start_epoch = 0
    resume_value = cfg.get("train", {}).get("resume_checkpoint")
    resume_path = resolve_resume_checkpoint(ckpt_dir, resume_value)
    if resume_path is not None:
        start_epoch, global_step, best = load_training_checkpoint(
            resume_path, model, optimizer=optimizer, velocity_field=velocity_field, teacher=teacher
        )
        if is_main_process():
            print(f"Resuming from epoch={start_epoch}, step={global_step}, best={best:.6f}")
    elif resume_value and is_main_process():
        print(f"WARNING: requested resume={resume_value!r}, but no usable checkpoint was found in {ckpt_dir}")
    epochs = int(cfg["train"].get("epochs", 100))
    steps_per_epoch = int(cfg["train"].get("steps_per_epoch", len(train_loader)))
    if start_epoch >= epochs and is_main_process():
        print(f"Checkpoint epoch {start_epoch} is already >= configured epochs {epochs}; nothing to train.")
    for epoch in range(start_epoch, epochs):
        if hasattr(train_loader.sampler, "set_epoch"):
            train_loader.sampler.set_epoch(epoch)
        model.train()
        if velocity_field is not None:
            velocity_field.train()
        iterator = iter(train_loader)
        pbar = range(steps_per_epoch)
        if is_main_process():
            pbar = tqdm(pbar, desc=f"epoch {epoch+1}/{epochs}", dynamic_ncols=True)
        for _ in pbar:
            try:
                batch = next(iterator)
            except StopIteration:
                iterator = iter(train_loader)
                batch = next(iterator)
            batch = move_to_device(batch, device)
            optimizer.zero_grad(set_to_none=True)
            with torch.cuda.amp.autocast(enabled=bool(cfg.get("amp", True)) and device.type == "cuda"):
                if mode in ("source_train", "oracle_train"):
                    loss, logs = supervised_step(model, batch, cfg)
                elif mode in ("self_train", "peft"):
                    loss, logs = pseudo_step(model, teacher, batch, cfg)
                elif mode == "proto_align":
                    loss, logs = proto_align_step(model, teacher, batch, memory, cfg)
                elif mode == "sacflow_fm":
                    loss, logs = sacflow_forward_step(model, teacher, velocity_field, batch, memory, cfg)
                else:
                    raise ValueError(f"Unknown train mode {mode}")
            scaler.scale(loss).backward()
            if float(cfg["optim"].get("grad_clip_norm", 0) or 0) > 0:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_([p for p in params if p.requires_grad], float(cfg["optim"].get("grad_clip_norm")))
            scaler.step(optimizer)
            scaler.update()
            if mode in ("self_train", "peft", "sacflow_fm", "proto_align"):
                update_ema(teacher, unwrap_model(model), float(cfg["train"].get("ema_decay", 0.995)))
            global_step += 1
            red_logs = {}
            for k, v in logs.items():
                if torch.is_tensor(v):
                    red_logs[f"train/{k}"] = float(reduce_mean(v.float()).item())
                else:
                    red_logs[f"train/{k}"] = v
            if is_main_process() and global_step % int(cfg["train"].get("log_every", 20)) == 0:
                red_logs["train/epoch"] = epoch + 1
                red_logs["train/lr"] = optimizer.param_groups[0]["lr"]
                wandb_log(wandb_run, red_logs, step=global_step)
                if hasattr(pbar, "set_postfix"):
                    pbar.set_postfix({"loss": f"{red_logs.get('train/loss_total', 0):.4f}", "step": global_step})
        if val_loader is not None and ((epoch + 1) % int(cfg["train"].get("val_every", 1)) == 0):
            metrics = evaluate(unwrap_model(model), val_loader, cfg, device)
            if is_main_process():
                score = metrics.get("val/dice_mean", -1e9)
                print(f"Epoch {epoch+1} validation: {metrics}")
                wandb_log(wandb_run, metrics, step=global_step)
                if score > best:
                    best = score
                    save_checkpoint(
                        ckpt_dir / "best.pt", model, optimizer, epoch+1, global_step, best,
                        velocity_field, cfg, teacher=teacher,
                        include_optimizer=bool(cfg["train"].get("save_optimizer_in_best", False)),
                    )
            # Keep all ranks synchronized after rank0 best-checkpoint writing.
            barrier()
        # Always keep a resumable last.pt. Do not create epoch_N.pt files unless explicitly requested.
        save_checkpoint(
            ckpt_dir / "last.pt", model, optimizer, epoch+1, global_step, best,
            velocity_field, cfg, teacher=teacher, include_optimizer=True,
        )
        # Ensure rank0 has finished writing last.pt before other ranks start the next epoch.
        barrier()
        if bool(cfg["train"].get("keep_epoch_checkpoints", False)) and int(cfg["train"].get("save_every", 0) or 0) > 0:
            if (epoch + 1) % int(cfg["train"].get("save_every", 5)) == 0:
                save_checkpoint(
                    ckpt_dir / f"epoch_{epoch+1}.pt", model, optimizer, epoch+1, global_step, best,
                    velocity_field, cfg, teacher=teacher, include_optimizer=True,
                )
                barrier()
    # final last.pt is already saved after every epoch; save once more for completeness.
    save_checkpoint(ckpt_dir / "last.pt", model, optimizer, epochs, global_step, best, velocity_field, cfg, teacher=teacher, include_optimizer=True)
    return unwrap_model(model)