scripts/pretrain.py

#!/usr/bin/env python
"""Pretrain MuViTMAE3d from a YAML config file."""

import argparse
import math
import os
import signal
from argparse import Namespace
from collections import defaultdict
from pathlib import Path

import numpy as np
import yaml
import torch
from torch.utils.data import DataLoader, Dataset
from torch.utils.data._utils.collate import default_collate

from miao import VolumeDataset, load_config
from muvit.mae import MuViTMAE3d

torch.set_float32_matmul_precision("medium")


MODALITY_PREFIXES = {
    "flyem": ("flyem", "flyem fib-sem", "flyem cns", "flyem_fibsem"),
    "fafb": ("fafb", "fafb v14", "fafb_sstem"),
    "flyliconn": ("flyliconn",),
    "nisb_base_em": ("nisb_base", "nisb base", "base_em", "base em"),
    "liconn": ("nisb_liconn", "liconn"),
}


def modality_from_volume_name(name: str) -> str:
    text = name.lower().replace("-", " ")
    for modality, prefixes in MODALITY_PREFIXES.items():
        if any(text.startswith(prefix) or prefix in text for prefix in prefixes):
            return modality
    raise ValueError(f"Could not infer modality from volume name {name!r}")


def config_modalities(miao_cfg) -> list[str]:
    return sorted({modality_from_volume_name(volume.name) for volume in miao_cfg.volumes})


def collate_allow_none(batch):
    """Default-collate dict samples while allowing image-only label=None."""
    if batch and isinstance(batch[0], dict):
        out = {}
        for key in batch[0]:
            values = [item[key] for item in batch]
            if all(value is None for value in values):
                if key == "label":
                    out[key] = torch.empty((len(values), 0), dtype=torch.long)
                else:
                    out[key] = None
            else:
                out[key] = default_collate(values)
        return out
    return default_collate(batch)


def identity_collate(batch):
    """Return already-collated batch objects unchanged for batch-yielding datasets."""
    return batch


def _per_sample_p5_p95(x: torch.Tensor) -> torch.Tensor:
    flat = x.flatten(start_dim=1)
    q = torch.quantile(flat.float(), torch.tensor([0.05, 0.95], device=x.device), dim=1)
    lo = q[0].view(-1, *([1] * (x.ndim - 1)))
    hi = q[1].view(-1, *([1] * (x.ndim - 1)))
    y = x.float().clamp(lo, hi)
    return (y - lo) / (hi - lo).clamp_min(1e-6)


def _match_one_to_reference(sample: torch.Tensor, reference: torch.Tensor) -> torch.Tensor:
    flat = sample.flatten().float()
    order = torch.argsort(flat)
    ref_sorted = torch.sort(reference.flatten().float()).values
    if ref_sorted.numel() != flat.numel():
        pos = torch.linspace(0, ref_sorted.numel() - 1, flat.numel(), device=flat.device)
        lo = pos.floor().long()
        hi = pos.ceil().long()
        alpha = pos - lo.float()
        ref_sorted = ref_sorted[lo] * (1 - alpha) + ref_sorted[hi] * alpha
    out = torch.empty_like(flat)
    out[order] = ref_sorted
    return out.view_as(sample)


def _batch_hist_match(x: torch.Tensor, modalities: list[str]) -> torch.Tensor:
    by_modality: dict[str, list[int]] = defaultdict(list)
    for idx, modality in enumerate(modalities):
        by_modality[modality].append(idx)
    if not by_modality or min(len(v) for v in by_modality.values()) == 0:
        raise ValueError(f"Pipeline B requires at least one sample per modality, got {modalities}")
    per_modality_count = min(len(v) for v in by_modality.values())
    refs = []
    for modality in sorted(by_modality):
        for idx in by_modality[modality][:per_modality_count]:
            refs.append(x[idx].float().flatten())
    reference = torch.cat(refs)
    matched = torch.stack([_match_one_to_reference(sample, reference) for sample in x])
    lo = matched.flatten(start_dim=1).min(dim=1).values.view(-1, *([1] * (x.ndim - 1)))
    hi = matched.flatten(start_dim=1).max(dim=1).values.view(-1, *([1] * (x.ndim - 1)))
    return (matched - lo) / (hi - lo).clamp_min(1e-6)


def _extract_modalities(batch: dict) -> list[str]:
    meta = batch.get("meta", {})
    volumes = meta.get("volume") if isinstance(meta, dict) else None
    if isinstance(volumes, str):
        volumes = [volumes]
    if volumes is None:
        raise ValueError("Preprocessing requires batch['meta']['volume']")
    return [modality_from_volume_name(str(volume)) for volume in volumes]


def apply_input_preprocessing(batch: dict, pipeline: str) -> dict:
    if pipeline in ("none", "disabled", None):
        return batch
    x = batch["img"]
    if pipeline == "pipeline_a":
        batch["img"] = _per_sample_p5_p95(x)
        return batch
    if pipeline == "pipeline_b":
        batch["img"] = _batch_hist_match(x, _extract_modalities(batch))
        return batch
    raise ValueError(f"Unknown preprocessing pipeline {pipeline!r}")


class PreprocessingCollate:
    """Collate normal sample lists and apply the configured input preprocessing."""

    def __init__(self, pipeline: str) -> None:
        self.pipeline = pipeline

    def __call__(self, samples: list[dict]) -> dict:
        return apply_input_preprocessing(collate_allow_none(samples), self.pipeline)


class BalancedBatchDataset(Dataset):
    """Batch-yielding wrapper that enforces exact per-modality batch balance.

    miao.VolumeDataset samples its source volume inside __getitem__, so a
    PyTorch Sampler cannot force a balanced modality mix. This wrapper draws
    samples until each modality quota is filled, collates them, and applies
    preprocessing. Validation can be made deterministic by seeding numpy from
    the batch index before each balanced draw.
    """

    def __init__(
        self,
        dataset: VolumeDataset,
        *,
        batch_size: int,
        modalities: list[str],
        pipeline: str,
        deterministic: bool = False,
        seed: int = 20260622,
        attempt_limit: int = 20000,
    ) -> None:
        if not modalities:
            raise ValueError("BalancedBatchDataset requires at least one modality")
        if batch_size % len(modalities) != 0:
            raise ValueError(
                f"batch_size={batch_size} must be divisible by num_modalities={len(modalities)} "
                f"for exact balanced batches ({modalities})"
            )
        self.dataset = dataset
        self.batch_size = int(batch_size)
        self.modalities = sorted(modalities)
        self.per_modality = self.batch_size // len(self.modalities)
        self.pipeline = pipeline
        self.deterministic = bool(deterministic)
        self.seed = int(seed)
        self.attempt_limit = int(attempt_limit)

    def __len__(self) -> int:
        return max(1, math.ceil(len(self.dataset) / self.batch_size))

    def __getitem__(self, idx: int) -> dict:
        if self.deterministic:
            np.random.seed(self.seed + int(idx))
        buckets: dict[str, list[dict]] = {modality: [] for modality in self.modalities}
        attempts = 0
        while any(len(items) < self.per_modality for items in buckets.values()):
            sample = self.dataset[attempts % max(1, len(self.dataset))]
            modality = modality_from_volume_name(sample["meta"]["volume"])
            if modality in buckets and len(buckets[modality]) < self.per_modality:
                buckets[modality].append(sample)
            attempts += 1
            if attempts > self.attempt_limit:
                counts = {key: len(value) for key, value in buckets.items()}
                raise RuntimeError(
                    f"Could not build balanced batch after {self.attempt_limit} draws; "
                    f"modalities={self.modalities}, counts={counts}"
                )
        samples = []
        for modality in self.modalities:
            samples.extend(buckets[modality])
        return apply_input_preprocessing(collate_allow_none(samples), self.pipeline)


def make_pretrain_loader(
    ds: VolumeDataset,
    ds_cfg,
    train_cfg: dict,
    *,
    batch_size_key: str,
    shuffle: bool,
    deterministic: bool = False,
    seed: int = 20260622,
) -> DataLoader:
    pipeline = train_cfg.get("preprocessing_pipeline", "none")
    balanced = bool(train_cfg.get("preprocessing_balanced_batches", False)) or pipeline == "pipeline_b"
    num_workers = int(train_cfg["data_loader_workers"])
    prefetch_factor = train_cfg.get("data_loader_prefetch_factor", 2)
    batch_size = int(train_cfg[batch_size_key])
    common = dict(
        num_workers=num_workers,
        pin_memory=True,
        persistent_workers=num_workers > 0,
        prefetch_factor=prefetch_factor if num_workers > 0 else None,
        multiprocessing_context="spawn" if num_workers > 0 else None,
    )
    if balanced:
        batch_ds = BalancedBatchDataset(
            ds,
            batch_size=batch_size,
            modalities=config_modalities(ds_cfg),
            pipeline=pipeline,
            deterministic=deterministic,
            seed=seed,
            attempt_limit=int(train_cfg.get("preprocessing_balance_attempt_limit", 20000)),
        )
        return DataLoader(batch_ds, batch_size=None, shuffle=False, collate_fn=identity_collate, **common)
    return DataLoader(
        ds,
        batch_size=batch_size,
        shuffle=shuffle,
        collate_fn=PreprocessingCollate(pipeline),
        **common,
    )


def main():
    parser = argparse.ArgumentParser(description="MuViT MAE 3D Pretraining")
    parser.add_argument("-m", "--main_config", type=str, help="Path to YAML config file which contains training and model hyperparameters")
    parser.add_argument("-t", "--train_config", type=str, help="Path to YAML config file containing training dataset parameters")
    parser.add_argument("-v", "--val_config", type=str, nargs="+", help="One or more YAML config files containing validation dataset parameters")
    parser.add_argument("--val-name", type=str, nargs="*", default=None, help="Optional validation loader names matching --val_config")
    parser.add_argument("-r", "--run-name", type=str, default=None, help="W&B run name (overrides config)")
    parser.add_argument("--output", type=str, default=None, help="Output directory override")
    parser.add_argument("--resume-ckpt", type=str, default=None, help="Optional Lightning checkpoint path to resume trainer state")
    args = parser.parse_args()

    with open(args.main_config) as f:
        cfg = yaml.safe_load(f)

    model_cfg = cfg["model"]
    train_cfg = cfg["training"]
    log_cfg = cfg["logging"]

    print("=" * 60)
    print("Configuration")
    print("=" * 60)
    print(yaml.dump(cfg, default_flow_style=False, sort_keys=False))
    print("=" * 60)

    train_ds_cfg = load_config(args.train_config)
    val_config_paths = list(args.val_config or [])
    if not val_config_paths:
        raise ValueError("At least one --val_config is required")
    val_names = args.val_name or [Path(p).stem for p in val_config_paths]
    if len(val_names) != len(val_config_paths):
        raise ValueError("--val-name count must match --val_config count")
    val_ds_cfgs = [load_config(p) for p in val_config_paths]

    # --- Datasets ---
    train_ds = VolumeDataset(train_ds_cfg)
    val_datasets = [VolumeDataset(cfg) for cfg in val_ds_cfgs]

    # --- DataLoaders ---
    train_dl = make_pretrain_loader(
        train_ds,
        train_ds_cfg,
        train_cfg,
        batch_size_key="batch_size",
        shuffle=True,
    )
    val_dls = [
        make_pretrain_loader(
            val_ds,
            val_cfg,
            train_cfg,
            batch_size_key="val_batch_size",
            shuffle=False,
            deterministic=True,
            seed=int(train_cfg.get("preprocessing_val_seed", 20260622)) + 1000 * idx,
        )
        for idx, (val_ds, val_cfg) in enumerate(zip(val_datasets, val_ds_cfgs))
    ]

    # --- Model ---
    if model_cfg.get("pretrained_folder"):
        print(f"Loading MuViT MAE checkpoint from {model_cfg['pretrained_folder']}")
        model = MuViTMAE3d.from_folder(model_cfg["pretrained_folder"])
    elif model_cfg.get("pretrained_hf"):
        print(f"Loading MuViT MAE checkpoint from Hugging Face {model_cfg['pretrained_hf']}")
        model = MuViTMAE3d.from_hf(model_cfg["pretrained_hf"])
    else:
        model = MuViTMAE3d(
            levels=tuple(model_cfg.get("levels", (1, 2, 4))),
            patch_size=model_cfg["patch_size"],
            num_layers=model_cfg["num_layers"],
            dim=model_cfg["dim"],
            num_layers_decoder=model_cfg["num_layers_decoder"],
            dim_decoder=model_cfg["dim_decoder"],
            heads=model_cfg["heads"],
            masking_ratio=model_cfg["masking_ratio"],
            decoder_mode=model_cfg["decoder_mode"],
            rotary_mode=model_cfg["rotary_mode"],
            loss=model_cfg["loss"],
            masking_mode=model_cfg.get("masking_mode", "dirichlet"),
            use_level_embed=model_cfg.get("use_level_embed", True),
            attention_mode=model_cfg.get("attention_mode", "all"),
            dropout=model_cfg.get("dropout", 0.0),
        )

    # --- Flatten config for hyperparameter logging ---
    flat_cfg = {}
    for section in (model_cfg, train_cfg, log_cfg):
        flat_cfg.update(section)
    flat_cfg.update(
        {
            "main_config": args.main_config,
            "train_config": args.train_config,
            "val_config": val_config_paths,
            "val_names": val_names,
            "resume_ckpt": args.resume_ckpt,
        }
    )

    # --- Training ---
    output = Path(args.output or log_cfg["output"])
    output.mkdir(parents=True, exist_ok=True)

    # Handle SIGTERM (e.g. SLURM job cancellation) gracefully
    def _sigterm_handler(signum, frame):
        raise KeyboardInterrupt("Received SIGTERM")

    signal.signal(signal.SIGTERM, _sigterm_handler)

    try:
        fit_kw = dict(
            train_dataloader=train_dl,
            val_dataloader=val_dls if len(val_dls) > 1 else val_dls[0],
            output=output,
            num_epochs=train_cfg["num_epochs"],
            min_epochs=train_cfg.get("min_epochs", 0),
            lr=train_cfg["lr"],
            warmup_epochs=train_cfg["warmup_epochs"],
            logger=log_cfg["logger"],
            run_name=args.run_name or log_cfg["run_name"],
            wandb_project=os.environ.get("WANDB_PROJECT", log_cfg["wandb_project"]),
            wandb_entity=os.environ.get("WANDB_ENTITY", log_cfg.get("wandb_entity")),
            wandb_tags=log_cfg.get("wandb_tags"),
            precision=train_cfg["precision"],
            gradient_clip_val=train_cfg["gradient_clip_val"],
            strategy=train_cfg["strategy"],
            num_nodes=train_cfg["num_nodes"],
            args_namespace=Namespace(**flat_cfg),
            check_val_every_n_epoch=train_cfg.get("check_val_every_n_epoch", 100),
            early_stopping_monitor=train_cfg.get("early_stopping_monitor", "mae/val/mean_loss"),
            early_stopping_mode=train_cfg.get("early_stopping_mode", "min"),
            early_stopping_patience=train_cfg.get("early_stopping_patience"),
            early_stopping_min_delta=train_cfg.get("early_stopping_min_delta", 0.0),
            checkpoint_save_top_k=train_cfg.get("checkpoint_save_top_k", 3),
            checkpoint_save_last=train_cfg.get("checkpoint_save_last", True),
            val_names=val_names,
            encoder_diagnostics_every_n_val_epochs=train_cfg.get("encoder_diagnostics_every_n_val_epochs", 0),
            max_encoder_diagnostic_batches=train_cfg.get("max_encoder_diagnostic_batches", 16),
            alignment_weight=train_cfg.get("alignment_weight", 0.0),
            mmd_weight=train_cfg.get("mmd_weight", 0.0),
            mmd_log_only=train_cfg.get("mmd_log_only", False),
            mmd_feature_mode=train_cfg.get("mmd_feature_mode", "token"),
            mmd_max_features_per_domain=train_cfg.get("mmd_max_features_per_domain", 1024),
            mmd_gamma=train_cfg.get("mmd_gamma"),
            mmd_normalize_features=train_cfg.get("mmd_normalize_features", True),
            style_aug_enabled=train_cfg.get("style_aug_enabled", False),
            style_aug_apply_to_mae=train_cfg.get("style_aug_apply_to_mae", False),
            style_intensity_scale_min=train_cfg.get("style_intensity_scale_min", 0.8),
            style_intensity_scale_max=train_cfg.get("style_intensity_scale_max", 1.2),
            style_intensity_shift=train_cfg.get("style_intensity_shift", 0.1),
            style_gamma_min=train_cfg.get("style_gamma_min", 0.8),
            style_gamma_max=train_cfg.get("style_gamma_max", 1.25),
            style_noise_std=train_cfg.get("style_noise_std", 0.03),
            style_blur_prob=train_cfg.get("style_blur_prob", 0.25),
            style_z_dropout_prob=train_cfg.get("style_z_dropout_prob", 0.0),
            vicreg_weight=train_cfg.get("vicreg_weight", 0.0),
            vicreg_mode=train_cfg.get("vicreg_mode", "none"),
            vicreg_lambda_inv=train_cfg.get("vicreg_lambda_inv", 25.0),
            vicreg_lambda_var=train_cfg.get("vicreg_lambda_var", 25.0),
            vicreg_lambda_cov=train_cfg.get("vicreg_lambda_cov", 1.0),
            vicreg_gamma=train_cfg.get("vicreg_gamma", 1.0),
            vicreg_projector_hidden_dim=train_cfg.get("vicreg_projector_hidden_dim", 1024),
            vicreg_projector_output_dim=train_cfg.get("vicreg_projector_output_dim", 512),
            vicreg_token_max_features=train_cfg.get("vicreg_token_max_features", 4096),
            vicreg_all_gather=train_cfg.get("vicreg_all_gather", True),
            style_consistency_weight=train_cfg.get("style_consistency_weight", 0.0),
            style_consistency_mode=train_cfg.get("style_consistency_mode", "token"),
            style_consistency_loss=train_cfg.get("style_consistency_loss", "cosine"),
            style_consistency_token_max_features=train_cfg.get("style_consistency_token_max_features", 4096),
            style_consistency_projector_hidden_dim=train_cfg.get("style_consistency_projector_hidden_dim", 1024),
            style_consistency_projector_output_dim=train_cfg.get("style_consistency_projector_output_dim", 512),
            style_consistency_detach_target=train_cfg.get("style_consistency_detach_target", False),
            structure_weight=train_cfg.get("structure_weight", 0.0),
            structure_targets=train_cfg.get("structure_targets", "affinity_x,affinity_y,affinity_z,boundary"),
            structure_level_idx=train_cfg.get("structure_level_idx", 0),
            structure_background_label=train_cfg.get("structure_background_label", 0),
            structure_affinity_weight=train_cfg.get("structure_affinity_weight", 1.0),
            structure_boundary_weight=train_cfg.get("structure_boundary_weight", 1.0),
            structure_head_hidden_dim=train_cfg.get("structure_head_hidden_dim", 0),
            structure_head_dropout=train_cfg.get("structure_head_dropout", 0.0),
            token_barlow_weight=train_cfg.get("token_barlow_weight", 0.0),
            token_barlow_lambda_offdiag=train_cfg.get("token_barlow_lambda_offdiag", 0.005),
            token_barlow_level_idx=train_cfg.get("token_barlow_level_idx"),
            token_barlow_max_features=train_cfg.get("token_barlow_max_features", 4096),
            token_barlow_projector_hidden_dim=train_cfg.get("token_barlow_projector_hidden_dim", 1024),
            token_barlow_projector_output_dim=train_cfg.get("token_barlow_projector_output_dim", 512),
            token_barlow_all_gather=train_cfg.get("token_barlow_all_gather", True),
            token_barlow_structure_stratified=train_cfg.get("token_barlow_structure_stratified", True),
            token_barlow_boundary_fraction=train_cfg.get("token_barlow_boundary_fraction", 0.5),
            ckpt_path=args.resume_ckpt,
        )
        if train_cfg.get("accelerator") is not None:
            fit_kw["accelerator"] = train_cfg["accelerator"]
        if train_cfg.get("devices") is not None:
            fit_kw["devices"] = train_cfg["devices"]
        model.fit(**fit_kw)
    except KeyboardInterrupt:
        print("\nTraining interrupted.")
    finally:
        # Clean up GPU memory
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.reset_peak_memory_stats()
        # Finish wandb run cleanly
        import wandb
        if wandb.run is not None:
            wandb.finish()


if __name__ == "__main__":
    main()