train_multimae.py

"""
MultiMAE3D Pretraining Script with Cross-Modal Prediction + Anatomy-Aware Masking

Usage:
    # Single GPU
    python train_multimae.py --batch_size 4

    # Multi-GPU DDP
    torchrun --nproc_per_node=8 train_multimae.py --batch_size 4

    # With cross-modal prediction + anatomy-aware masking
    torchrun --nproc_per_node=8 train_multimae.py --batch_size 4 \
        --enable_cross_modal --use_anatomy_masking --atlas_path altas/AAL116_standard.nii.gz
"""

import os
import argparse
import math
import time
import numpy as np
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from tensorboardX import SummaryWriter
from tqdm import tqdm

from models.multimae3d import create_multimae3d
from pretrain_dataloader_v2 import MultiModalPretrainDataset
from anatomy_masking import (
    AnatomyAwareMasking,
    create_ema_teacher,
    update_ema_teacher,
    extract_teacher_attention,
)


# =============================================================================
# Distributed setup
# =============================================================================

def setup_distributed():
    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
        rank = int(os.environ["RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
        local_rank = int(os.environ["LOCAL_RANK"])
        dist.init_process_group(backend="nccl")
        torch.cuda.set_device(local_rank)
        return rank, world_size, local_rank
    return 0, 1, 0


def cleanup_distributed():
    if dist.is_initialized():
        try:
            torch.cuda.synchronize()
            dist.barrier()
            dist.destroy_process_group()
        except Exception as e:
            print(f"[Rank {dist.get_rank()}] Warning: cleanup failed: {e}")
            try:
                dist.destroy_process_group()
            except Exception:
                pass


# =============================================================================
# Training loop
# =============================================================================

def cosine_ema_momentum(epoch, total_epochs, start=0.996, end=1.0):
    """Cosine schedule for EMA momentum: start → end over training."""
    progress = epoch / max(total_epochs, 1)
    return end - (end - start) * (1 + math.cos(math.pi * progress)) / 2


def cross_modal_lambda_schedule(epoch, warmup_epochs, target_lambda):
    """λ=0 for first warmup_epochs, then linear ramp to target over next warmup_epochs."""
    if epoch <= warmup_epochs:
        return 0.0
    ramp_progress = min(1.0, (epoch - warmup_epochs) / max(warmup_epochs, 1))
    return target_lambda * ramp_progress


def train_one_epoch(
    model, dataloader, optimizer, epoch, writer,
    rank=0, device="cuda", global_step=0, grad_clip=0.5,
    enable_cross_modal=False, cross_modal_lambda=0.1,
    cross_modal_warmup_epochs=10, total_epochs=1200,
    ema_momentum_start=0.996, ema_momentum_end=1.0,
    anatomy_masking=None, anatomy_ema_teacher=None,
):
    model.train()
    model_inner = model.module if hasattr(model, "module") else model

    # Compute schedules for this epoch
    ema_momentum = cosine_ema_momentum(
        epoch, total_epochs, ema_momentum_start, ema_momentum_end,
    )
    effective_lambda = cross_modal_lambda_schedule(
        epoch, cross_modal_warmup_epochs, cross_modal_lambda,
    ) if enable_cross_modal else 0.0

    # Anatomy masking: compute mask probabilities for this epoch
    mask_probs = None
    if anatomy_masking is not None:
        mask_probs = anatomy_masking.get_mask_probs(epoch, total_epochs)
        if mask_probs is not None:
            mask_probs = mask_probs.to(device)

    use_dynamic_anatomy = (
        anatomy_masking is not None
        and anatomy_ema_teacher is not None
        and anatomy_masking.importance_mode in ('dynamic', 'combined')
    )

    total_loss = 0.0
    total_cross_loss = 0.0
    per_mod_losses = {name: 0.0 for name in ["T1", "T2", "Flair", "PET"]}
    num_batches = 0

    pbar = tqdm(dataloader, desc=f"Epoch {epoch}", disable=(rank != 0))

    for batch_idx, batch in enumerate(pbar):
        images = batch["images"].to(device)      # [B, 4, 128, 128, 128]
        observed = batch["observed"].to(device)   # [B, 4]

        # Forward with anatomy-aware masking
        output = model(images, observed, return_loss=True, patch_mask_probs=mask_probs)
        mae_loss = output["loss"]
        cross_loss = output.get("cross_modal_loss", torch.tensor(0.0, device=device))

        # Combined loss
        loss = mae_loss + effective_lambda * cross_loss

        # Backward
        optimizer.zero_grad()
        loss.backward()

        # Gradient clipping
        grad_norm = nn.utils.clip_grad_norm_(
            [p for p in model.parameters() if p.requires_grad],
            max_norm=grad_clip,
        )

        optimizer.step()

        # EMA update of cross-modal teacher (after optimizer step)
        if enable_cross_modal:
            model_inner.update_teacher(ema_momentum)

        # EMA update of anatomy masking teacher (every step)
        if anatomy_ema_teacher is not None:
            update_ema_teacher(anatomy_ema_teacher, model, momentum=anatomy_masking.ema_momentum)

        # Periodically extract teacher attention and update dynamic importance
        iteration = global_step + batch_idx
        if use_dynamic_anatomy and iteration > 0 and iteration % anatomy_masking.attention_update_freq == 0:
            tb = min(anatomy_masking.teacher_batch_size, images.shape[0])
            with torch.no_grad():
                attn = extract_teacher_attention(
                    anatomy_ema_teacher, images[:tb], observed[:tb],
                )
            anatomy_masking.update_dynamic_importance(attn)
            # Recompute mask probs with updated dynamic importance
            new_probs = anatomy_masking.get_mask_probs(epoch, total_epochs)
            if new_probs is not None:
                mask_probs = new_probs.to(device)

        # Logging
        mae_val = mae_loss.item()
        cross_val = cross_loss.item() if torch.is_tensor(cross_loss) else cross_loss
        combined_val = loss.item()
        total_loss += mae_val
        total_cross_loss += cross_val
        num_batches += 1

        for name, mod_loss in output["per_modality_loss"].items():
            per_mod_losses[name] += mod_loss.item()

        if rank == 0 and writer is not None:
            step = global_step + batch_idx
            writer.add_scalar("Train/Batch/MAE_Loss", mae_val, step)
            writer.add_scalar("Train/Batch/Total_Loss", combined_val, step)
            if enable_cross_modal:
                writer.add_scalar("Train/Batch/Cross_Modal_Loss", cross_val, step)
                writer.add_scalar("Train/Batch/Cross_Lambda", effective_lambda, step)
                writer.add_scalar("Train/Batch/EMA_Momentum", ema_momentum, step)
            writer.add_scalar("Train/Batch/Grad_Norm", grad_norm.item(), step)
            writer.add_scalar("Train/Batch/LR", optimizer.param_groups[0]["lr"], step)
            for name, mr in output["mask_ratios"].items():
                writer.add_scalar(f"Train/Batch/MaskRatio_{name}", mr, step)

        if rank == 0:
            postfix = {"mae": f"{mae_val:.4f}"}
            if enable_cross_modal and effective_lambda > 0:
                postfix["cross"] = f"{cross_val:.4f}"
            pbar.set_postfix(postfix)

    avg_loss = total_loss / max(num_batches, 1)
    avg_cross_loss = total_cross_loss / max(num_batches, 1)
    avg_mod_losses = {k: v / max(num_batches, 1) for k, v in per_mod_losses.items()}

    if rank == 0 and writer is not None:
        writer.add_scalar("Train/Epoch/MAE_Loss", avg_loss, epoch)
        writer.add_scalar("Train/Epoch/Total_Loss",
                          avg_loss + effective_lambda * avg_cross_loss, epoch)
        if enable_cross_modal:
            writer.add_scalar("Train/Epoch/Cross_Modal_Loss", avg_cross_loss, epoch)
        for name, ml in avg_mod_losses.items():
            writer.add_scalar(f"Train/Epoch/Loss_{name}", ml, epoch)
        writer.add_scalar("Train/Epoch/LR", optimizer.param_groups[0]["lr"], epoch)

        # Log anatomy masking curriculum info
        if anatomy_masking is not None:
            info = anatomy_masking.get_curriculum_info(epoch, total_epochs)
            writer.add_scalar("Anatomy/Phase", info['phase'], epoch)
            writer.add_scalar("Anatomy/Temperature", info['temperature'], epoch)
            if 'importance_max' in info:
                writer.add_scalar("Anatomy/Importance_Max", info['importance_max'], epoch)
                writer.add_scalar("Anatomy/Importance_Min", info['importance_min'], epoch)
                writer.add_scalar("Anatomy/Importance_Mean", info['importance_mean'], epoch)
            if 'prob_ratio' in info:
                writer.add_scalar("Anatomy/Prob_MaxMinRatio", info['prob_ratio'], epoch)

    return avg_loss, avg_mod_losses, num_batches


# =============================================================================
# Checkpoint management
# =============================================================================

def save_checkpoint(
    model, optimizer, scheduler, epoch, loss, best_loss,
    global_step, save_dir, rank=0, save_freq=100,
    anatomy_masking=None, anatomy_ema_teacher=None,
):
    if rank != 0:
        return

    os.makedirs(save_dir, exist_ok=True)
    model_inner = model.module if hasattr(model, "module") else model

    # Always save latest (for resume)
    resume_ckpt = {
        "epoch": epoch,
        "model_state_dict": model_inner.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "scheduler_state_dict": scheduler.state_dict() if scheduler else None,
        "loss": loss,
        "best_loss": best_loss,
        "global_step": global_step,
    }
    if anatomy_masking is not None:
        resume_ckpt["anatomy_masking_state"] = anatomy_masking.state_dict()
    if anatomy_ema_teacher is not None:
        ema_inner = anatomy_ema_teacher.module if hasattr(anatomy_ema_teacher, "module") else anatomy_ema_teacher
        resume_ckpt["anatomy_ema_teacher_state_dict"] = ema_inner.state_dict()
    torch.save(resume_ckpt, os.path.join(save_dir, "latest.pth"))

    # Periodic save (encoder only, for downstream)
    if epoch % save_freq == 0:
        encoder_ckpt = {
            "epoch": epoch,
            "encoder_state_dict": {
                k: v for k, v in model_inner.state_dict().items()
                if k.startswith("encoder.") or k.startswith("input_adapters.") or k.startswith("pos_embed") or k.startswith("global_tokens")
            },
            "loss": loss,
        }
        torch.save(encoder_ckpt, os.path.join(save_dir, f"encoder_epoch_{epoch}.pth"))

    # Save best
    if loss <= best_loss:
        encoder_ckpt = {
            "epoch": epoch,
            "encoder_state_dict": {
                k: v for k, v in model_inner.state_dict().items()
                if k.startswith("encoder.") or k.startswith("input_adapters.") or k.startswith("pos_embed") or k.startswith("global_tokens")
            },
            "full_model_state_dict": model_inner.state_dict(),
            "loss": loss,
        }
        torch.save(encoder_ckpt, os.path.join(save_dir, "best_model.pth"))


# =============================================================================
# Main
# =============================================================================

def main():
    parser = argparse.ArgumentParser(description="MultiMAE3D Pretraining")

    # Data
    parser.add_argument("--excel_dir", type=str,
                        default="./data/Match_data_path/pretraining_processed")
    parser.add_argument("--batch_size", type=int, default=4, help="Per-GPU batch size")
    parser.add_argument("--num_workers", type=int, default=8)
    parser.add_argument("--augmentation", action="store_true", default=True)
    parser.add_argument("--no_augmentation", action="store_false", dest="augmentation")

    # Model
    parser.add_argument("--img_size", type=int, default=128)
    parser.add_argument("--patch_size", type=int, default=16)
    parser.add_argument("--embed_dim", type=int, default=768)
    parser.add_argument("--depth", type=int, default=12)
    parser.add_argument("--num_heads", type=int, default=12)
    parser.add_argument("--decoder_embed_dim", type=int, default=384)
    parser.add_argument("--decoder_depth", type=int, default=2)
    parser.add_argument("--decoder_num_heads", type=int, default=12)
    parser.add_argument("--mask_ratio", type=float, default=0.75)
    parser.add_argument("--use_dirichlet", action="store_true", default=True)
    parser.add_argument("--no_dirichlet", action="store_false", dest="use_dirichlet")
    parser.add_argument("--dirichlet_alpha", type=float, default=1.0)
    parser.add_argument("--drop_path_rate", type=float, default=0.0)

    # Cross-modal mutual prediction
    parser.add_argument("--enable_cross_modal", action="store_true", default=False,
                        help="Enable cross-level mutual prediction (MRI↔PET)")
    parser.add_argument("--cross_modal_lambda", type=float, default=0.1,
                        help="Weight for cross-modal loss (search: 0.01-1.0)")
    parser.add_argument("--cross_modal_warmup_epochs", type=int, default=10,
                        help="Epochs with λ=0 before linear ramp")
    parser.add_argument("--ema_momentum_start", type=float, default=0.996,
                        help="EMA momentum at start of training (for cross-modal teacher)")
    parser.add_argument("--ema_momentum_end", type=float, default=1.0,
                        help="EMA momentum at end of training (for cross-modal teacher)")

    # Anatomy-aware masking
    parser.add_argument("--use_anatomy_masking", action="store_true", default=False,
                        help="Enable anatomy-aware adaptive masking")
    parser.add_argument("--atlas_path", type=str, default="altas/AAL116_standard.nii.gz",
                        help="Path to brain atlas NIfTI file (128^3, in data space)")
    parser.add_argument("--importance_mode", type=str, default="combined",
                        choices=["static", "dynamic", "combined"],
                        help="Importance scoring mode: static (AD prior), dynamic (EMA attention), combined")
    parser.add_argument("--anatomy_w_high", type=float, default=3.0,
                        help="Importance weight for AD-critical regions")
    parser.add_argument("--anatomy_w_mid", type=float, default=1.5,
                        help="Importance weight for other gray matter regions")
    parser.add_argument("--anatomy_w_low", type=float, default=0.3,
                        help="Importance weight for non-brain patches")
    parser.add_argument("--anatomy_temp_target", type=float, default=1.0,
                        help="Target temperature for masking softmax (lower = more focused)")
    parser.add_argument("--anatomy_temp_start", type=float, default=5.0,
                        help="Starting temperature at Phase 2 onset")
    parser.add_argument("--anatomy_phase1_end", type=float, default=0.2,
                        help="End of Phase 1 (uniform masking) as fraction of total epochs")
    parser.add_argument("--anatomy_phase2_end", type=float, default=0.7,
                        help="End of Phase 2 (transition) as fraction of total epochs")
    parser.add_argument("--anatomy_ema_momentum", type=float, default=0.998,
                        help="EMA momentum for anatomy masking teacher model")
    parser.add_argument("--attention_update_freq", type=int, default=200,
                        help="Extract teacher attention every N iterations")
    parser.add_argument("--teacher_batch_size", type=int, default=2,
                        help="Batch size for teacher attention extraction")
    parser.add_argument("--dynamic_weight", type=float, default=0.5,
                        help="Weight of dynamic importance in combined mode [0, 1]")

    # Training
    parser.add_argument("--epochs", type=int, default=1200)
    parser.add_argument("--warmup_epochs", type=int, default=40)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--weight_decay", type=float, default=0.05)
    parser.add_argument("--grad_clip", type=float, default=0.5)
    parser.add_argument("--seed", type=int, default=42)

    # Save
    parser.add_argument("--save_dir", type=str, default="./pretrain_checkpoints/multimae")
    parser.add_argument("--save_freq", type=int, default=100)
    parser.add_argument("--log_dir", type=str, default="./logs/multimae")
    parser.add_argument("--resume", type=str, default="", help="Path to latest.pth for resume (restores epoch/optimizer)")
    parser.add_argument("--pretrain_weights", type=str, default="",
                        help="Path to pretrained checkpoint — loads model weights only, starts from epoch 1")

    args = parser.parse_args()

    # Seed
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    # Distributed
    rank, world_size, local_rank = setup_distributed()
    device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")

    if rank == 0:
        print("=" * 70)
        print("MultiMAE3D Pretraining")
        if args.enable_cross_modal:
            print("  + Cross-Modal Mutual Prediction ENABLED")
        if args.use_anatomy_masking:
            print("  + Anatomy-Aware Adaptive Masking ENABLED")
            print(f"    Atlas: {args.atlas_path}")
            print(f"    Mode: {args.importance_mode}")
            print(f"    Weights: high={args.anatomy_w_high}, mid={args.anatomy_w_mid}, low={args.anatomy_w_low}")
            print(f"    Temperature: {args.anatomy_temp_start} -> {args.anatomy_temp_target}")
            print(f"    Curriculum: Phase1 end={args.anatomy_phase1_end}, Phase2 end={args.anatomy_phase2_end}")
        print("=" * 70)
        print(f"World size: {world_size}, Device: {device}")
        print(f"Config: {vars(args)}")
        print("=" * 70)

    # Dataset
    dataset = MultiModalPretrainDataset(
        excel_dir=args.excel_dir,
        image_size=(args.img_size, args.img_size, args.img_size),
        augmentation=args.augmentation,
        modality_dropout_prob=0.0,  # No dropout — natural missing is enough
        min_modalities=1,
    )

    if world_size > 1:
        sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
        dataloader = torch.utils.data.DataLoader(
            dataset, batch_size=args.batch_size, sampler=sampler,
            num_workers=args.num_workers, pin_memory=True, drop_last=True,
        )
    else:
        dataloader = torch.utils.data.DataLoader(
            dataset, batch_size=args.batch_size, shuffle=True,
            num_workers=args.num_workers, pin_memory=True, drop_last=True,
        )

    if rank == 0:
        print(f"Dataset: {len(dataset)} samples, {len(dataloader)} batches/epoch")

    # Model
    model = create_multimae3d(
        img_size=args.img_size,
        patch_size=args.patch_size,
        embed_dim=args.embed_dim,
        depth=args.depth,
        num_heads=args.num_heads,
        decoder_embed_dim=args.decoder_embed_dim,
        decoder_depth=args.decoder_depth,
        decoder_num_heads=args.decoder_num_heads,
        mask_ratio=args.mask_ratio,
        use_dirichlet=args.use_dirichlet,
        dirichlet_alpha=args.dirichlet_alpha,
        drop_path_rate=args.drop_path_rate,
        enable_cross_modal=args.enable_cross_modal,
    )
    model = model.to(device)

    if rank == 0:
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        print(f"Model params: {total_params:,} total, {trainable_params:,} trainable ({trainable_params/1e6:.1f}M)")

    # Anatomy-aware masking setup (before DDP wrapping)
    anatomy_masking = None
    anatomy_ema_teacher = None

    if args.use_anatomy_masking:
        anatomy_masking = AnatomyAwareMasking(
            img_size=args.img_size,
            patch_size=args.patch_size,
            atlas_path=args.atlas_path,
            w_high=args.anatomy_w_high,
            w_mid=args.anatomy_w_mid,
            w_low=args.anatomy_w_low,
            temperature_target=args.anatomy_temp_target,
            temperature_start=args.anatomy_temp_start,
            phase1_end=args.anatomy_phase1_end,
            phase2_end=args.anatomy_phase2_end,
            ema_momentum=args.anatomy_ema_momentum,
            attention_update_freq=args.attention_update_freq,
            teacher_batch_size=args.teacher_batch_size,
            importance_mode=args.importance_mode,
            dynamic_weight=args.dynamic_weight,
        )

        if rank == 0:
            info = anatomy_masking.get_curriculum_info(1, args.epochs)
            print(f"Anatomy masking initialized: {anatomy_masking.num_patches} patches, "
                  f"importance range [{info.get('importance_min', 'N/A'):.3f}, "
                  f"{info.get('importance_max', 'N/A'):.3f}]")

        # Create EMA teacher for dynamic importance (before DDP)
        if args.importance_mode in ('dynamic', 'combined'):
            anatomy_ema_teacher = create_ema_teacher(model)
            if rank == 0:
                print(f"Anatomy EMA teacher created (momentum={args.anatomy_ema_momentum})")

    # DDP wrapping (after EMA teacher creation)
    if world_size > 1:
        model = DDP(model, device_ids=[local_rank], find_unused_parameters=True)

    # Optimizer (only trainable params — excludes teacher EMA parameters)
    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=args.lr, weight_decay=args.weight_decay,
    )

    # LR Scheduler: linear warmup + cosine annealing
    def lr_lambda(epoch):
        if epoch < args.warmup_epochs:
            return (epoch + 1) / args.warmup_epochs
        else:
            progress = (epoch - args.warmup_epochs) / max(args.epochs - args.warmup_epochs, 1)
            return 0.5 * (1.0 + np.cos(np.pi * progress))

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

    # TensorBoard
    writer = None
    if rank == 0:
        log_dir = os.path.join(args.log_dir, f"seed_{args.seed}")
        os.makedirs(log_dir, exist_ok=True)
        writer = SummaryWriter(log_dir)
        print(f"TensorBoard: {log_dir}")

    # Resume
    start_epoch = 1
    best_loss = float("inf")
    global_step = 0

    if args.pretrain_weights and os.path.isfile(args.pretrain_weights):
        # Load model weights only — epoch/optimizer/scheduler stay fresh (start from epoch 1)
        if rank == 0:
            print(f"Loading pretrained weights: {args.pretrain_weights}")
        ckpt = torch.load(args.pretrain_weights, map_location=f"cuda:{local_rank}")
        model_inner = model.module if hasattr(model, "module") else model
        state_dict = ckpt.get("model_state_dict", ckpt.get("full_model_state_dict", ckpt))
        missing, unexpected = model_inner.load_state_dict(state_dict, strict=False)
        if rank == 0:
            print(f"  Loaded weights from epoch {ckpt.get('epoch', '?')}")
            if missing:
                print(f"  Missing keys ({len(missing)}): {missing[:5]}{'...' if len(missing) > 5 else ''}")
            if unexpected:
                print(f"  Unexpected keys ({len(unexpected)}): {unexpected[:5]}{'...' if len(unexpected) > 5 else ''}")
        # Initialize cross-modal teacher from the loaded student weights
        if args.enable_cross_modal:
            model_inner.init_teacher_from_student()
            if rank == 0:
                print("  Cross-modal teacher initialized from loaded student weights")
        # Initialize anatomy EMA teacher from loaded weights
        if anatomy_ema_teacher is not None:
            anatomy_ema_teacher.load_state_dict(state_dict, strict=False)
            if rank == 0:
                print("  Anatomy EMA teacher initialized from pretrained weights")
        if rank == 0:
            print(f"  Training from epoch 1 with fresh optimizer/scheduler")
        del ckpt

    elif args.resume and os.path.isfile(args.resume):
        # Full resume: restore model + optimizer + scheduler + epoch counter
        if rank == 0:
            print(f"Resuming from: {args.resume}")
        ckpt = torch.load(args.resume, map_location=f"cuda:{local_rank}")
        model_inner = model.module if hasattr(model, "module") else model
        missing, unexpected = model_inner.load_state_dict(ckpt["model_state_dict"], strict=False)
        if rank == 0 and (missing or unexpected):
            print(f"  load_state_dict: {len(missing)} missing, {len(unexpected)} unexpected keys")
        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
        if ckpt.get("scheduler_state_dict"):
            scheduler.load_state_dict(ckpt["scheduler_state_dict"])
        start_epoch = ckpt["epoch"] + 1
        best_loss = ckpt.get("best_loss", float("inf"))
        global_step = ckpt.get("global_step", 0)
        if args.enable_cross_modal:
            model_inner.init_teacher_from_student()
            if rank == 0:
                print("  Cross-modal teacher re-initialized from loaded student weights")
        # Restore anatomy masking state
        if anatomy_masking is not None and "anatomy_masking_state" in ckpt:
            anatomy_masking.load_state_dict(ckpt["anatomy_masking_state"])
            if rank == 0:
                print("  Restored anatomy masking state")
        # Restore anatomy EMA teacher
        if anatomy_ema_teacher is not None and "anatomy_ema_teacher_state_dict" in ckpt:
            anatomy_ema_teacher.load_state_dict(ckpt["anatomy_ema_teacher_state_dict"])
            if rank == 0:
                print("  Restored anatomy EMA teacher state")
        if rank == 0:
            print(f"Resumed from epoch {ckpt['epoch']}, best_loss={best_loss:.4f}")
        del ckpt

    # Training loop
    for epoch in range(start_epoch, args.epochs + 1):
        if world_size > 1:
            dataloader.sampler.set_epoch(epoch)

        t0 = time.time()
        avg_loss, avg_mod_losses, num_batches = train_one_epoch(
            model, dataloader, optimizer, epoch, writer,
            rank=rank, device=device, global_step=global_step,
            grad_clip=args.grad_clip,
            enable_cross_modal=args.enable_cross_modal,
            cross_modal_lambda=args.cross_modal_lambda,
            cross_modal_warmup_epochs=args.cross_modal_warmup_epochs,
            total_epochs=args.epochs,
            ema_momentum_start=args.ema_momentum_start,
            ema_momentum_end=args.ema_momentum_end,
            anatomy_masking=anatomy_masking,
            anatomy_ema_teacher=anatomy_ema_teacher,
        )
        global_step += num_batches
        elapsed = time.time() - t0

        # Step LR scheduler
        scheduler.step()

        if rank == 0:
            mod_str = ", ".join(f"{k}: {v:.4f}" for k, v in avg_mod_losses.items() if v > 0)
            phase_str = ""
            if anatomy_masking is not None:
                info = anatomy_masking.get_curriculum_info(epoch, args.epochs)
                phase_str = f" | Phase {info['phase']}, tau={info['temperature']:.2f}"
            print(f"Epoch {epoch}/{args.epochs} | Loss: {avg_loss:.4f} | {mod_str}{phase_str} | Time: {elapsed:.1f}s")

        # Save checkpoint
        save_checkpoint(
            model, optimizer, scheduler, epoch, avg_loss, best_loss,
            global_step, args.save_dir, rank, args.save_freq,
            anatomy_masking=anatomy_masking, anatomy_ema_teacher=anatomy_ema_teacher,
        )

        if avg_loss < best_loss:
            best_loss = avg_loss
            if rank == 0:
                print(f"  -> New best loss: {best_loss:.4f}")

    # Cleanup
    if world_size > 1 and dist.is_initialized():
        torch.cuda.synchronize()
        dist.barrier()

    cleanup_distributed()
    if writer is not None:
        writer.add_scalar("Final/Best_Loss", best_loss, 0)
        writer.close()

    if rank == 0:
        print(f"\nTraining complete! Best loss: {best_loss:.4f}")


if __name__ == "__main__":
    main()