#!/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()