diff --git a/dino-pretraining/README.md b/dino-pretraining/README.md index 0cecd9e..60600c4 100644 --- a/dino-pretraining/README.md +++ b/dino-pretraining/README.md @@ -1,22 +1,38 @@ -# Atlas-space 3-D DINOv3 for MR-RATE +# Co-registered 3-D DINOv3 for MR-RATE This module adapts FORA's CT DINOv3 training strategy to MR-RATE's -**atlas-registered multi-sequence MRI**. It is self-supervised: reports and -pathology labels are not used. +**co-registered multi-sequence MRI** (`coreg_space`, default; `atlas_space` remains +selectable). It is self-supervised: reports and pathology labels are not used for +training; they are used only by the validation probe (below) to score checkpoints. ## Data and view contract -Production training reads the existing MR-RATE-atlas NIfTI tree directly: +Production training reads an extracted MR-RATE-coreg NIfTI tree directly: ```text -/batchXX//atlas_img/*.nii.gz +/batchXX//coreg_img/*.nii.gz # --space coreg_space (default) +/batchXX//atlas_img/*.nii.gz # --space atlas_space ``` +**Why coreg, not atlas.** Both spaces use the same rigid (ANTs, Mattes MI) registration, so +both are voxel-aligned across sequences, which the cross-sequence views need. They differ +in the target grid: +- **Coreg** keeps every sequence on the study's T1w center scan's native grid (e.g. + 0.47×0.47×1.2 mm, 190×240×240 mm field of view). +- **Atlas** resamples into the 1 mm MNI box (193×229×193 mm). + +The model input is 1.0×0.5×0.5 mm, so from atlas the 0.5 mm in-plane detail is +interpolation, and the MNI box crops the skull base and neck (−12% head volume on a +checked study), which matters for the spinal labels. `native_space` is rejected +(its sequences are not aligned). + It imports the canonical discovery and per-volume preprocessing functions from `contrastive-pretraining/scripts/data.py`, the loader used by previous MR-RATE MIL training. Therefore atlas selection, canonical RAS orientation, physical resampling, z-score normalization, posterior shift, and crop/padding are exactly -shared rather than reimplemented. `--space atlas_space` maps to `atlas_img`. +shared rather than reimplemented. `--space coreg_space` maps to `coreg_img`, +`--space atlas_space` to `atlas_img`. Checkpoints record their space, and resume refuses +to cross spaces (checkpoints written before the coreg switch are atlas-space). An optional preprocessed volume cache remains available with `--preprocessed-dir`. This is separate from cached MIL: cached MIL stores the @@ -34,9 +50,9 @@ keeps only its current study stack in memory, avoiding repeated NIfTI decoding. For each anchor: 1. Global view 1 uses the anchor sequence. -2. Global view 2 uses a different atlas-aligned sequence with probability 0.75 - (or the anchor when no second sequence exists). The two crops overlap by at - least 75% on every axis. +2. Global view 2 uses a different atlas-aligned sequence with probability 0.25 + (`--cross-sequence-probability`; otherwise, or when no second sequence exists, + the anchor). The two crops overlap by at least 25% on every axis. 3. Local crops stay inside the global intersection and may come from any aligned sequence. 4. DINO and KoLeo therefore learn anatomy shared across MR contrasts. @@ -53,15 +69,55 @@ masking, distributed KoLeo, optional Gram anchoring, EMA teacher, atomic full state checkpoints, exact sampler/RNG resume, FSDP2, BF16, optional H200 FP8, and compile caches on node-local `/tmp`. +## Recipe fixes ported from FORA CT-DINO (Sep 2026) + +The CT pipeline this module was adapted from started with the same DINOv3-7B continuation +settings and failed; each item below is a diagnosed failure and the fix the working FORA +ViT-L run uses (code: `mr_dino/recipe.py`, `objective.py`, `data.py`; tests: +`tests/test_fora_recipe.py`). All are on by default. + +| Problem seen in FORA | Fix (default) | +|---|---| +| DINOv3 plain-linear prototype layer (std 0.02) never sharpened at lr 5e-5 → flat Sinkhorn targets for 10k steps (iBOT target entropy 8.0/11.5 nats, top-1 0.006) | cosine prototypes on both heads (`--normalize-prototypes`) | +| 1k warmup + 7B settings drove positional drift in the last blocks | DINOv2 ViT-L recipe: lr 1e-3 at batch 1024 (sqrt-scaled), warmup 2,500, teacher-temperature warmup 0.04→0.07 over 10k, weight decay 0.04→0.4, EMA 0.992→1, AdamW β2 0.999, clip 3, layer-wise lr decay 0.9, patch-embed lr ×0.2, prototype layers frozen for 1,250 steps | +| 7B continuation was worse than a from-scratch ViT-L at the same budget | train from scratch at a trainable size: FORA used `--arch large` (ViT-L); MR uses `--arch hplus` (ViT-H+, 0.84B). 16k/16k prototypes, heads 2048/256, no FP8 | +| global crops overlapping ≥75% were near-duplicates | ≥25% overlap per axis (`--global-overlap`), locals inside one of the two globals | +| head-facing block encoded where a token sat in its crop (last-block position leakage 0.20, crop-quadrant prototype fields) | second global snapped to the patch lattice + **cross-crop twin-token objective** (`--cross-view-weight 1.0`): the student's patch prediction in one crop matches the teacher target of the same voxels in the other crop; leakage fell to 0.09 and stayed flat | +| patch targets balanced over all masked tokens at once | Sinkhorn within 2×2×2 spatial bins of the crop (`--position-bins`, CAPI) | +| raw loss hid whether the student learned (CE = H + KL) | `metrics.jsonl` logs `dino/ibot_target_entropy`, `_top1`, `_prototype_usage`, `dino_kl`, `ibot_kl`, `cross`, `cross_kl`, `cross_pairs` | + +Already fixed here before the port: the distributed Sinkhorn uses the actual global token count. + +**MR-specific choice:** twin tokens are paired only when both global crops show the same +sequence (`--cross-view-sequences same`), preserving this module's rule that cross-sequence +images are never patch-level targets. With `--cross-sequence-probability 0.25` that is ~75% of +samples; `any` would also pair co-registered voxels across contrasts (contrast-invariant patch +features, at the risk of suppressing lesions visible in only one sequence). + +**Cross-sequence global pairs default to 0.25 (was 0.75).** A DINO target shared by a T1 and a +FLAIR crop rewards discarding contrast-specific findings (FLAIR-only gliosis, SWI-only bleeds), +and in an aligned space the cheapest shared signal is the crop's position. Many "different" +sequences are the same contrast in another plane, so the true cross-contrast share is lower still. +A small non-zero rate keeps one embedding space across contrasts for MIL bags that mix sequences; +0 vs 0.25 is to be settled with the validation MIL probe. + +Production: `scripts/train_mrdino.sbatch` (ViT-H+, 4 nodes, 3 phases; see below). +`train_32n_7b.sbatch` / `train_32n_hplus.sbatch` keep the old recipe for reference only. + +**Healthy early training (FORA reference, ViT-L, global batch 128–1024):** iBOT target top-1 +rises from ~0.003 to >0.05 within 2.5k steps, patch target entropy falls, prototype usage stays +≳0.4, `cross_kl` ~1.5–2.5 without a rising trend. Gram anchoring (stage `gram`) gave no gain in +FORA and stays optional. + ## Optional input-volume cache From `contrastive-pretraining/`: ```bash python scripts/preprocess_volumes.py \ - --data_folder /path/to/MR-RATE-atlas/mri \ + --data_folder /path/to/MR-RATE-coreg/mri \ --out_dir /path/to/mrrate_preprocessed \ - --space atlas_space \ + --space coreg_space \ --normalizer zscore \ --num_workers 8 ``` @@ -104,49 +160,148 @@ sbatch scripts/smoke_dummy.sbatch The job succeeds only after `step_00000002/COMPLETE` exists. -## Production training - -The production launcher uses the existing MR-RATE-atlas dataset directly: +## Production training (3 phases, ViT-H+, 4 nodes) + +`scripts/train_mrdino.sbatch` runs every phase. Set `STAGE`; everything else has a phase +default. It reads the extracted MR-RATE-coreg tree directly (`SPACE=atlas_space` switches to +the atlas tree). + +**Model: ViT-H+ (0.84B).** This is the official DINOv3 H+ shape with 3-D patch embedding +and RoPE: 32 blocks, width 1280, 20 heads, SwiGLU ×6, 2×16×16 patch = 2×8×8 mm tokens. The +heads are 2048 → 256 with 16,384 cosine prototypes each. +- **Why not 7B:** the 7B (the previous MR default) failed in FORA with the old recipe, and + 4 nodes cannot train it. +- **Why not ViT-L:** ViT-L (0.30B, `ARCH=large`) is the size proven with this recipe in + FORA CT, but it is smaller than wanted. +- **Learning rate:** it was tuned on ViT-L. Watch the target metrics in the first + 2–3k steps (see "Healthy early training" above). + +| | Phase 1 `pretrain` | Phase 2 `gram` | Phase 3 `highres` | +|---|---|---|---| +| steps | **50,000** (~10 passes over ~635k sequences) | 10,000 | 5,000 | +| starts from | scratch | phase-1 end | phase-2 end (or phase-1 end, see below) | +| global crops | 2 × 64×192×192 vox (64×96×96 mm, 4,608 tokens) | same | 2 × 64×384×384 (whole axial plane, 18,432 tokens) | +| local crops | 8 × 32×96×96 | 8 × 32×96×96 | 4 × 32×192×192 | +| batch / GPU (global on 4 nodes) | 8 (128) | 8 (128) | 4 (64) | +| activation checkpointing | on (selective) | on | on | +| peak lr (at batch 1024, sqrt-scaled) | 1e-3 → 3.5e-4 | 5e-4 → 1.8e-4 | 4e-4 → 1e-4 | +| warmup / end lr | 2,500 / cosine to 1e-6 at step end | 250 / same | 500 / same | +| weight decay | 0.04 → 0.4 | 0.1 → 0.4 | 0.04 → 0.2 | +| teacher temperature | 0.04 → 0.07 over 10k | 0.07 fixed | 0.07 fixed | +| teacher EMA | 0.992 → 1 | 0.996 → 1 | 0.996 → 1 | +| prototype freeze | first 1,250 steps | none | none | +| Gram weight | 0 | 1.0 | 1.5 | + +Phases 2 and 3 use FORA CT-DINO's tested settings (`phase2_gram_recipe.args`, +`phase3_hires_recipe.args`). A phase restarts its schedule at step 0, so reusing phase-1 +values would redo the warmups: teacher temperature back to 0.04, prototypes frozen again. + +**Settings shared by all phases:** +- **Losses:** DINO 1.0 + iBOT 1.0 (block masks 10–50%) + twin-token 1.0 + KoLeo 0.1. + Sinkhorn targets use 2×2×2 spatial bins. +- **Sequences:** cross-sequence global pairs 25%, twin tokens within the same sequence. +- **Optimiser:** AdamW β2 0.999, clip 3, layer-wise lr decay 0.9, patch-embed lr ×0.2, BF16. +- **Infrastructure:** FSDP2, compiled blocks, checkpoint every 500 steps (last 6 kept). + +**Memory and speed per H200** (1-node benchmark, H+, compiled, real crop sizes): + +| Configuration | Result | +|---|---| +| phase-1 crops, batch 8, no checkpointing | **out of memory** | +| phase-1 crops, batch 8 + checkpointing | 0.39 steps/s, 32 GiB (default) | +| phase-1 crops, batch 4, no checkpointing | 0.95 steps/s, 83 GiB | +| highres crops, batch 2, no checkpointing | 128 GiB (too tight) | +| highres crops, batch 2 + checkpointing | 0.23 steps/s, 43 GiB | +| highres crops, batch 4 + checkpointing | 0.10 steps/s, 81 GiB (default) | + +So on 4 nodes, as long as data loading keeps up, expect roughly: +- phase 1: ~36 h; +- phase 2: ~10 h (the Gram anchor adds a forward pass); +- phase 3: ~14 h. + +Each phase is longer than one 24 h job; the launcher checkpoints at the wall-time signal +and requeues itself. ```bash cd dino-pretraining -DATA_FOLDER=/path/to/MR-RATE-atlas/mri \ -SPLITS_CSV=/path/to/splits.csv \ -OUTPUT=/path/to/mrdino3d_7b/pretrain \ -sbatch scripts/train_32n_7b.sbatch -``` +export DATA_FOLDER=/path/to/MR-RATE-coreg/mri SPLITS_CSV=/path/to/splits.csv +R=/path/to/mrdino3d_hplus -Important overrides are `STEPS`, `BATCH_SIZE`, `GRAD_ACCUM_STEPS`, `WORKERS`, -`LOCAL_CROPS`, `CROSS_SEQUENCE_PROBABILITY`, `WARMUP_STEPS`, `RESUME`, `FP8`, -and `COMPILE`. The default warmup is 1,000 optimizer steps. `RESUME=latest` -loads the latest checkpoint only when its `COMPLETE` marker exists. +# phase 1 (+ validation probe of chosen steps, run alongside on the login node) +STAGE=pretrain OUTPUT=$R/pretrain sbatch scripts/train_mrdino.sbatch +nohup scripts/probe_watch.sh $R/pretrain 5000 10000 20000 30000 40000 50000 & -For later stages: +# phase 2 from the phase-1 end +STAGE=gram RESUME=$R/pretrain/checkpoints/step_00050000 OUTPUT=$R/gram sbatch scripts/train_mrdino.sbatch +nohup scripts/probe_watch.sh $R/gram 5000 10000 & -```bash -STAGE=gram RESUME=/path/to/pretrain/checkpoints/step_00100000 \ -DATA_FOLDER=/path/to/MR-RATE-atlas/mri \ -OUTPUT=/path/to/mrdino3d_7b/gram \ -sbatch scripts/train_32n_7b.sbatch +# phase 3 from the phase-2 end, only if the gram probe is >= the phase-1 probe; +# otherwise RESUME=$R/pretrain/checkpoints/step_00050000 (FORA CT: Gram gave no gain) +STAGE=highres RESUME=$R/gram/checkpoints/step_00010000 OUTPUT=$R/highres sbatch scripts/train_mrdino.sbatch +nohup scripts/probe_watch.sh $R/highres 2500 5000 & ``` -The launcher catches the Slurm wall-time signal, writes a complete distributed -checkpoint with per-rank sampler and RNG state, and requeues the same job. Do -not start production before `scripts/smoke_dummy.sbatch` passes in the same -container and DINOv3 checkout. - -The same 32-node H+ fallback used by CT-DINO is also available: +**Rules:** +- **`RESUME` for phases 2/3:** it names the previous phase's checkpoint only for the first + start. Once the phase has its own complete checkpoint, a requeued job continues from that + one. A phase-2/3 launch without `RESUME` and without its own checkpoints is refused. +- **Resume does not cross spaces or arches:** checkpoints record their space, and the + probe rebuilds the arch from the checkpoint metadata. +- **`STEPS` is the real run length:** the cosine decay ends there, so set it before + launching, not mid-run. +- **Overrides:** `ARCH`, `STEPS`, `BATCH_SIZE`, `ACT_CKPT`, `GRAD_ACCUM_STEPS`, `WORKERS`, + `LOCAL_CROPS`, `CROSS_SEQUENCE_PROBABILITY`, `CROSS_VIEW_SEQUENCES`, `GRAM_WEIGHT`, + `RESUME`, `SPACE`, `COMPILE`. Container and paths: `SIF`, `DINOV3_ROOT`, `EXTRA_PIP`, + `ZIG_ARCHIVE` (see `scripts/train_node.sh`). +- **Site-specific `#SBATCH` lines:** edit `--partition`, the `--exclude` bad-node list, + and `--mail-user` before running elsewhere. Slurm refuses unknown node names. +- **Run `scripts/smoke_dummy.sbatch` first** (1 node, ~2 min) in the same container and + DINOv3 checkout. It trains, resumes and runs the probe on synthetic coreg data. + +`train_32n_7b.sbatch` and `train_32n_hplus.sbatch` are the old 7B-continuation recipe +(32-node DDP/FSDP) and are kept for reference only. + +## Validation probe (checkpoint scoring) + +This is the FORA CT-DINO protocol: frozen features, MIL on the validation split only, and +patient-level cross-validation. Train and test are never touched. + +1. **Export** (`python -m mr_dino.probe export`): copies the EMA-teacher backbone out of a + DCP checkpoint, so checkpoint rotation cannot delete it. +2. **Extract** (`python -m mr_dino.probe extract`, one worker per GPU): + - Runs every sequence of every labelled `val` study (3,764 studies, 3,413 patients) + through the training preprocessing of the checkpoint's own space/spacing/shape. + - Cuts each volume into phase-1-size tiles (64×192×192 voxels) and keeps the last + block's normalised patch tokens. + - Average-pools the tokens 2×2×2 over head foreground and drops air. + - Writes one fp16 token file plus one index per worker to node-local `/tmp`. +3. **CV** (`python -m mr_dino.probe cv --fold k`): trains MR-RATE's NeuroVFM + `ClassifyThenAggregate` head per patient-level fold (5 folds, 8 epochs, batch 8, + lr 2e-3, pos_weight clipped to [1, 100]). One bag holds all sequences of a study. +4. **Summarize**: out-of-fold macro and per-class AUROC/AUPRC → `results.json`, + `per_class.csv`, `summary.txt`. + +Labels default to the 14 merged neuroradiology groups +(`contrastive-pretraining/scripts/eval_labels/splits_merged_majority`). Set +`LABELS_CSV`/`SPLITS_CSV` to use the 32/37-pathology labels instead. ```bash -DATA_FOLDER=/path/to/MR-RATE-atlas/mri \ -OUTPUT=/path/to/mrdino3d_hplus/pretrain \ -sbatch scripts/train_32n_hplus.sbatch +# one checkpoint, one node (extract -> 5 folds -> summary; ~1-2 h on 4 H200): +DATA_FOLDER=/path/to/MR-RATE-coreg/mri \ +sbatch scripts/probe_checkpoint.sbatch /path/to/pretrain/checkpoints/step_00005000 /path/to/pretrain/probe/step_5000 +# follow a running stage: export each listed step as it appears, then submit its probe: +DATA_FOLDER=/path/to/MR-RATE-coreg/mri nohup scripts/probe_watch.sh /path/to/pretrain 2500 5000 10000 15000 20000 & +# inside a running training allocation instead of a new job: +srun --overlap --jobid JOB -N1 -n1 -w NODE scripts/probe_node.sh CKPT OUT_DIR ``` -Both production launchers include the established bad-node exclusions, Slurm -mail notifications, 24-hour requeue policy, deadline checkpointing, node-local -compiler caches, InfiniBand/NCCL settings, and optional split filtering. The -7B path is the production target; H+ is a smaller DDP fallback. +Use the probe to: +- decide how long to run phase 1; +- check that phase 2 (Gram) does not lose AUROC before phase 3 starts from it; +- settle open recipe questions, e.g. `--cross-sequence-probability` 0 vs 0.25. + +Compare per-class rows, not just the macro AUROC: FLAIR-, SWI- and DWI-dependent groups +are what cross-sequence choices can hurt. ## Lightweight/debug run @@ -154,7 +309,7 @@ compiler caches, InfiniBand/NCCL settings, and optional split filtering. The ```bash torchrun --standalone --nproc-per-node=1 -m mr_dino.train_ddp \ - --data-folder /path/to/MR-RATE-atlas/mri \ + --data-folder /path/to/MR-RATE-coreg/mri \ --output-dir /tmp/mrdino_debug \ --arch tiny --steps 10 --workers 0 --local-crops 2 \ --prototypes 64 --head-hidden-dim 128 \ diff --git a/dino-pretraining/mr_dino/__init__.py b/dino-pretraining/mr_dino/__init__.py index eee5563..378ae02 100644 --- a/dino-pretraining/mr_dino/__init__.py +++ b/dino-pretraining/mr_dino/__init__.py @@ -1,5 +1,5 @@ -"""Volumetric DINOv3 training for atlas-registered MR-RATE studies.""" +"""Volumetric DINOv3 training for aligned (coreg/atlas) MR-RATE studies.""" -from .data import CropSpec, MRAtlasDINO3DDataset, collate_dino3d +from .data import CropSpec, MRAlignedDINO3DDataset, collate_dino3d -__all__ = ["CropSpec", "MRAtlasDINO3DDataset", "collate_dino3d"] +__all__ = ["CropSpec", "MRAlignedDINO3DDataset", "collate_dino3d"] diff --git a/dino-pretraining/mr_dino/data.py b/dino-pretraining/mr_dino/data.py index 696f5b9..f018d3c 100644 --- a/dino-pretraining/mr_dino/data.py +++ b/dino-pretraining/mr_dino/data.py @@ -1,6 +1,10 @@ -"""Atlas-registered multi-sequence MR sampling for volumetric DINOv3. +"""Aligned multi-sequence MR sampling for volumetric DINOv3. -The default path discovers and preprocesses atlas NIfTIs with the canonical +Training uses MR-RATE's co-registered space by default (``coreg_space``: every +sequence rigidly registered to the study's T1w center scan and kept on that +scan's native grid, so no MNI field-of-view crop and no 1 mm resampling); +``atlas_space`` remains selectable. Both are voxel-aligned across sequences. +The default path discovers and preprocesses NIfTIs with the canonical loader used by the previous MR-RATE training code. The optional NPZ path reads the volume cache produced by that loader's ``preprocess_volumes.py`` helper. Global and local DINO views share a physical region, but may use different @@ -25,8 +29,12 @@ import torch.nn.functional as F from torch.utils.data import Dataset, Sampler +from .recipe import overlap_pairs, snap_second_start, token_offset + CACHE_MANIFEST_NAME = "_manifest.json" +ALIGNED_SPACES = ("coreg_space", "atlas_space") +DEFAULT_SPACE = "coreg_space" EXPECTED_CACHE_LAYOUT = "per_subject_stack" _PREVIOUS_DATA_MODULE = None @@ -209,21 +217,30 @@ def _cache_space_dir(preprocessed_dir: str, space: str) -> Path: return nested if nested.is_dir() else root -def validate_atlas_cache( +def check_aligned_space(space: str) -> str: + """Cross-sequence views need voxel alignment: native_space is never accepted.""" + if space not in ALIGNED_SPACES: + raise ValueError( + f"MR DINO requires an aligned space {ALIGNED_SPACES}; got {space!r} " + "(native_space sequences are not voxel-aligned)" + ) + return space + + +def validate_aligned_cache( preprocessed_dir: str, - space: str = "atlas_space", + space: str = DEFAULT_SPACE, expected_spacing: tuple[float, float, float] = (1.0, 0.5, 0.5), ) -> dict: - """Fail early unless this is a compatible MR-RATE atlas-space cache.""" - if space != "atlas_space": - raise ValueError("MR DINO requires space='atlas_space'; native/coreg inputs are not accepted") + """Fail early unless this is a compatible MR-RATE aligned-space cache.""" + check_aligned_space(space) space_dir = _cache_space_dir(preprocessed_dir, space) manifest_path = space_dir / CACHE_MANIFEST_NAME if not manifest_path.is_file(): raise FileNotFoundError(f"Missing MR-RATE preprocessing manifest: {manifest_path}") manifest = json.loads(manifest_path.read_text()) - if manifest.get("space") != "atlas_space": - raise ValueError(f"Cache is {manifest.get('space')!r}, expected 'atlas_space'") + if manifest.get("space") != space: + raise ValueError(f"Cache is {manifest.get('space')!r}, expected {space!r}") if manifest.get("layout") != EXPECTED_CACHE_LAYOUT: raise ValueError( f"Unsupported cache layout {manifest.get('layout')!r}; expected {EXPECTED_CACHE_LAYOUT!r}" @@ -234,13 +251,13 @@ def validate_atlas_cache( return manifest -def discover_atlas_cache( +def discover_aligned_cache( preprocessed_dir: str, splits_csv: str | None = None, split: str = "train", - space: str = "atlas_space", + space: str = DEFAULT_SPACE, ) -> list[dict]: - manifest = validate_atlas_cache(preprocessed_dir, space) + manifest = validate_aligned_cache(preprocessed_dir, space) space_dir = _cache_space_dir(preprocessed_dir, space) selected = _load_split_ids(splits_csv, split) samples = [] @@ -256,7 +273,7 @@ def discover_atlas_cache( "volume_shape": shape[1:], }) if not samples: - raise RuntimeError(f"No {split!r} atlas-registered studies found under {space_dir}") + raise RuntimeError(f"No {split!r} {space} studies found under {space_dir}") return samples @@ -306,14 +323,16 @@ def previous_mr_data_module(): return module -def discover_raw_atlas( +def discover_raw_aligned( data_folder: str, selected: set[str] | None = None, + space: str = DEFAULT_SPACE, ) -> list[dict]: - """Discover atlas NIfTIs through the previous training dataloader.""" + """Discover aligned NIfTIs (coreg_img/ or atlas_img/) through the previous training dataloader.""" + check_aligned_space(space) samples = [] previous = previous_mr_data_module() - for subject in previous.discover_subjects(data_folder, "atlas_space"): + for subject in previous.discover_subjects(data_folder, space): study_uid = str(subject["subject_id"]) if selected is not None and study_uid not in selected: continue @@ -324,17 +343,18 @@ def discover_raw_atlas( "n_sequences": len(paths), }) if not samples: - raise RuntimeError(f"No atlas-registered NIfTIs found under {data_folder}") + raise RuntimeError(f"No {space} NIfTIs found under {data_folder}") return samples -def discover_raw_atlas_split( +def discover_raw_aligned_split( data_folder: str, splits_csv: str | None = None, split: str = "train", + space: str = DEFAULT_SPACE, ) -> list[dict]: - """Discover a split using the previous loader's atlas-space convention.""" - return discover_raw_atlas(data_folder, _load_split_ids(splits_csv, split)) + """Discover a split using the previous loader's space convention.""" + return discover_raw_aligned(data_folder, _load_split_ids(splits_csv, split), space) def _raw_volume( @@ -356,8 +376,8 @@ def _raw_volume( return torch.from_numpy(np.ascontiguousarray(array)).to(torch.bfloat16) -class MRAtlasDINO3DDataset(Dataset): - """Study-level SSL samples from atlas-aligned, variable-count MR sequences.""" +class MRAlignedDINO3DDataset(Dataset): + """Study-level SSL samples from voxel-aligned, variable-count MR sequences.""" def __init__( self, @@ -366,23 +386,31 @@ def __init__( data_folder: str | None = None, splits_csv: str | None = None, split: str = "train", - space: str = "atlas_space", + space: str = DEFAULT_SPACE, crop_spec: CropSpec = CropSpec(), target_spacing: tuple[float, float, float] = (1.0, 0.5, 0.5), target_shape: tuple[int, int, int] = (256, 384, 384), posterior_shift_mm: float = 15.0, cache_files: list[str] | None = None, raw_samples: list[dict] | None = None, - cross_sequence_probability: float = 0.75, + cross_sequence_probability: float = 0.25, candidate_trials: int = 12, seed: int = 3407, + global_overlap: float = 0.25, + patch_size: tuple[int, int, int] = (2, 16, 16), ) -> None: if (preprocessed_dir is None) == (data_folder is None): raise ValueError("Pass exactly one of preprocessed_dir or data_folder") - if space != "atlas_space": - raise ValueError("Only atlas_space is valid for atlas-registered MR DINO") + self.space = check_aligned_space(space) if not 0 <= cross_sequence_probability <= 1: raise ValueError("cross_sequence_probability must be in [0, 1]") + if not 0 <= global_overlap <= 1: + raise ValueError("global_overlap must be in [0, 1]") + # FORA fix: the two global crops only need >= 25% overlap per axis (0.75 made them near + # duplicates), and the second crop is snapped to the patch lattice so that every token in + # the overlap has an exact twin for the cross-crop patch objective. + self.global_overlap = float(global_overlap) + self.patch_size = tuple(int(v) for v in patch_size) self.crop_spec = crop_spec self.target_shape = tuple(int(x) for x in target_shape) self.target_spacing = tuple(float(x) for x in target_spacing) @@ -395,10 +423,10 @@ def __init__( self._cached_names: list[str] | None = None self.preprocessed = preprocessed_dir is not None if self.preprocessed: - manifest = validate_atlas_cache(preprocessed_dir, space, target_spacing) + manifest = validate_aligned_cache(preprocessed_dir, space, target_spacing) self.target_shape = tuple(int(x) for x in manifest["target_shape"]) if cache_files is None: - self.samples = discover_atlas_cache(preprocessed_dir, splits_csv, split, space) + self.samples = discover_aligned_cache(preprocessed_dir, splits_csv, split, space) else: selected = _load_split_ids(splits_csv, split) self.samples = [ @@ -412,18 +440,18 @@ def __init__( if selected is None or Path(path).stem in selected ] if not self.samples: - raise RuntimeError("This rank received no atlas-registered cache studies") + raise RuntimeError(f"This rank received no {space} cache studies") else: selected = _load_split_ids(splits_csv, split) if raw_samples is None: - self.samples = discover_raw_atlas(data_folder, selected) + self.samples = discover_raw_aligned(data_folder, selected, space) else: self.samples = [ sample for sample in raw_samples if selected is None or sample["study_uid"] in selected ] if not self.samples: - raise RuntimeError("This rank received no atlas-registered NIfTI studies") + raise RuntimeError(f"This rank received no {space} NIfTI studies") for sample in self.samples: if "volume_shape" in sample and tuple(sample["volume_shape"]) != self.target_shape: raise ValueError( @@ -500,10 +528,13 @@ def __getitem__(self, key: int | tuple[int, int]) -> dict: first_start = _informative_start( stack[first_sequence], self.crop_spec.global_shape, rng, self.candidate_trials ) - global_starts = _overlapping_global_starts( - tuple(stack.shape[1:]), self.crop_spec.global_shape, rng, first=first_start + container = tuple(stack.shape[1:]) + first, second = _overlapping_global_starts( + container, self.crop_spec.global_shape, rng, first=first_start, + min_overlap=self.global_overlap, ) - lo, hi = _intersection_box(global_starts, self.crop_spec.global_shape) + second = snap_second_start(first, second, container, self.crop_spec.global_shape, self.patch_size) + global_starts = (first, second) teacher_globals, student_globals = [], [] for sequence, start in zip(global_sequences, global_starts): @@ -514,6 +545,11 @@ def __getitem__(self, key: int | tuple[int, int]) -> dict: local_views, local_sequences, local_starts = [], [], [] for _ in range(self.crop_spec.local_crops): sequence = rng.randrange(n_sequences) + # FORA fix: locals lie inside ONE of the globals (the intersection of two crops with + # 25% overlap can be thinner than a local crop). + parent = global_starts[rng.randrange(2)] + lo = parent + hi = tuple(min(p + g, c) for p, g, c in zip(parent, self.crop_spec.global_shape, container)) start = _contained_start(lo, hi, self.crop_spec.local_shape, rng) clean = _crop_at_start(stack[sequence], self.crop_spec.local_shape, start) local_views.append(_student_distortion(clean, rng).unsqueeze(0)) @@ -573,6 +609,7 @@ def collate_dino3d( samples: list[dict], patch_size: tuple[int, int, int] = (2, 16, 16), mask_ratio: tuple[float, float] = (0.1, 0.5), + cross_view_sequences: str = "same", ) -> dict: teacher = torch.stack([sample["teacher_global"] for sample in samples], dim=1) student = torch.stack([sample["student_global"] for sample in samples], dim=1) @@ -590,7 +627,27 @@ def collate_dino3d( masks = torch.stack(masks) indices = masks.flatten().nonzero().flatten() weights = (1.0 / masks.sum(-1).clamp(min=1)).unsqueeze(-1).expand_as(masks)[masks] + # Twin tokens of the two global crops (cross-crop patch objective). By default only when both + # globals show the SAME sequence, preserving this module's contract that cross-sequence images + # are never patch-level targets; "any" also pairs co-registered voxels across contrasts. + if cross_view_sequences not in {"same", "any", "none"}: + raise ValueError("cross_view_sequences must be 'same', 'any' or 'none'") + n = math.prod(grid) + pa, pb, ps = [], [], [] + for index, sample in enumerate(samples): + seq0, seq1 = sample["global_sequences"] + if cross_view_sequences == "none" or (cross_view_sequences == "same" and seq0 != seq1): + continue + g0, g1 = sample["global_starts"] + a, b = overlap_pairs(token_offset(g0, g1, patch_size), grid) + pa.append(a + index * n) + pb.append(b + index * n) + ps.append(torch.full((a.numel(),), index, dtype=torch.long)) + empty = torch.empty(0, dtype=torch.long) return { + "cross_a": torch.cat(pa) if pa else empty, + "cross_b": torch.cat(pb) if pb else empty, + "cross_sample": torch.cat(ps) if ps else empty, "teacher_global": teacher, "student_global": student, "student_local": local, diff --git a/dino-pretraining/mr_dino/objective.py b/dino-pretraining/mr_dino/objective.py index 917cd30..9680e4a 100644 --- a/dino-pretraining/mr_dino/objective.py +++ b/dino-pretraining/mr_dino/objective.py @@ -3,6 +3,7 @@ from __future__ import annotations import copy +import math from dataclasses import dataclass import torch @@ -12,6 +13,7 @@ from dinov3.layers.dino_head import DINOHead from .model import DinoVisionTransformer3D +from .recipe import binned_sinkhorn, entropy, normalize_prototype_layer, padded_head, position_bin_ids, target_stats class SSLNetwork3D(nn.Module): @@ -24,6 +26,7 @@ def __init__( ibot_hidden_dim: int = 4096, dino_bottleneck_dim: int = 384, ibot_bottleneck_dim: int = 384, + normalize_prototypes: bool = True, ) -> None: super().__init__() self.backbone = backbone @@ -41,6 +44,11 @@ def __init__( bottleneck_dim=ibot_bottleneck_dim, nlayers=3, ) + if normalize_prototypes: + # FORA fix: cosine prototypes let the teacher commit from step 0 (plain std-0.02 + # prototypes kept the Sinkhorn targets flat for thousands of steps). + normalize_prototype_layer(self.dino_head) + normalize_prototype_layer(self.ibot_head) def init_weights(self) -> None: self.backbone.init_weights() @@ -184,6 +192,9 @@ def __init__( ibot_loss_chunk_size: int = 1024, with_gram_anchor: bool = False, initialize_weights: bool = True, + normalize_prototypes: bool = True, + position_bins: tuple[int, int, int] | None = (2, 2, 2), + cross_view_weight: float = 1.0, ) -> None: super().__init__() self.student = SSLNetwork3D( @@ -194,7 +205,11 @@ def __init__( ibot_head_hidden_dim or head_hidden_dim, dino_bottleneck_dim, ibot_bottleneck_dim, + normalize_prototypes=normalize_prototypes, ) + self.patch_size = tuple(int(v) for v in backbone.patch_size) + self.position_bins = tuple(int(v) for v in position_bins) if position_bins else None + self.cross_view_weight = float(cross_view_weight) if initialize_weights: self.student.init_weights() self.teacher = copy.deepcopy(self.student).requires_grad_(False) @@ -248,6 +263,49 @@ def reset_gram_anchor_from_teacher(self) -> None: self.gram_anchor.load_state_dict(self.teacher.state_dict()) self.gram_anchor.eval() + def _cross_view(self, batch, t_patch, s_patch, teacher_temperature, grid, n_tokens, batch_size, weighted_mean): + """Cross-crop patch objective (FORA v14): for every token present in both global crops + (crops aligned on the patch lattice; ``cross_a``/``cross_b`` are flat indices of the twins + in view 0 / view 1), the student's patch prediction in one crop must match the teacher's + Sinkhorn target for the same voxels in the other crop, in both directions. Without it the + iBOT term only compares a token with the teacher's view of the SAME crop, and the + head-facing block learns where a token sits in the crop (position leakage).""" + device = s_patch.device + a = batch["cross_a"].to(device) + b = batch["cross_b"].to(device) + batch_size * n_tokens # view 1 follows view 0 + s_flat = s_patch.flatten(0, 1) + t_flat = t_patch.flatten(0, 1) + n_pairs = int(a.numel()) + total_pairs = torch.tensor(n_pairs, device=device) + if dist.is_initialized(): + dist.all_reduce(total_pairs) + zero = s_patch.sum() * 0.0 + if int(total_pairs.item()) == 0: # no sample on any rank has twins (all ranks agree) + return zero, {"cross": zero, "cross_kl": zero, "cross_pairs": zero} + with torch.no_grad(): + t_tokens = torch.cat([t_flat.index_select(0, b), t_flat.index_select(0, a)]) + t_logits, _ = padded_head(self.teacher.ibot_head, t_tokens) + teacher_rows = torch.cat([b, a]) + if self.position_bins is not None: + q = binned_sinkhorn(distributed_sinkhorn, t_logits, teacher_temperature, + position_bin_ids(teacher_rows, n_tokens, grid, self.position_bins), + math.prod(self.position_bins)) + else: + q = distributed_sinkhorn(t_logits, teacher_temperature) + s_tokens = torch.cat([s_flat.index_select(0, a), s_flat.index_select(0, b)]) + s_logits, anchor = padded_head(self.student.ibot_head, s_tokens) + if n_pairs == 0: + return zero + anchor, {"cross": zero, "cross_kl": zero, "cross_pairs": zero} + logp = F.log_softmax(s_logits.float() / self.student_temperature, dim=-1) + ce = -(q * logp).sum(-1) + sample_ids = torch.cat([batch["cross_sample"], batch["cross_sample"]]).to(device) + counts = torch.bincount(sample_ids, minlength=batch_size).clamp_min(1).float() + per_sample = torch.zeros(batch_size, device=device).scatter_add_(0, sample_ids, ce) / counts + per_sample_h = torch.zeros(batch_size, device=device).scatter_add_(0, sample_ids, entropy(q)) / counts + cross = weighted_mean(per_sample) + anchor + return cross, {"cross": cross.detach(), "cross_kl": (cross - weighted_mean(per_sample_h)).detach(), + "cross_pairs": torch.tensor(float(n_pairs), device=device)} + def forward(self, batch: dict, teacher_temperature: float, step: int) -> tuple[Tensor, dict[str, Tensor]]: teacher_images = batch["teacher_global"] # [2,B,1,D,H,W] student_images = batch["student_global"] @@ -284,8 +342,19 @@ def weighted_mean(per_sample: Tensor) -> Tensor: t_patch = tout["x_norm_patchtokens"] t_masked = t_patch.flatten(0, 1).index_select(0, mask_indices) t_patch_logits = self.teacher.ibot_head(t_masked) - t_cls_prob = distributed_sinkhorn(t_cls_logits, teacher_temperature).unflatten(0, (n_global, batch_size)) - t_patch_prob = distributed_sinkhorn(t_patch_logits, teacher_temperature) + t_cls_flat = distributed_sinkhorn(t_cls_logits, teacher_temperature) + t_cls_prob = t_cls_flat.unflatten(0, (n_global, batch_size)) + n_tokens = int(masks.shape[-1]) + grid = tuple(n // p for n, p in zip(teacher_images.shape[-3:], self.patch_size)) + if self.position_bins is not None: + # FORA fix (CAPI): balance patch targets within spatial bins of the crop. + t_patch_prob = binned_sinkhorn( + distributed_sinkhorn, t_patch_logits, teacher_temperature, + position_bin_ids(mask_indices, n_tokens, grid, self.position_bins), + math.prod(self.position_bins), + ) + else: + t_patch_prob = distributed_sinkhorn(t_patch_logits, teacher_temperature) del t_patch_logits, t_masked global_out, local_out = self.student.backbone( @@ -375,14 +444,43 @@ def weighted_mean(per_sample: Tensor) -> Tensor: ).unflatten(0, (n_global, batch_size)).mean(0) gram = weighted_mean(gram_each) + # Teacher-target diagnostics (FORA fix): CE = H(target) + KL, so the raw loss hides + # whether the student learns; log H, KL = CE - H (weighted like the losses), top-1, usage. + with torch.no_grad(): + cls_stats = target_stats(t_cls_flat) + patch_stats = target_stats(t_patch_prob) + h_cls = weighted_mean(entropy(t_cls_prob).mean(0)) + h_rows = entropy(t_patch_prob) + h_mask = torch.zeros(masks.shape[0], device=h_rows.device, dtype=h_rows.dtype) + h_mask.scatter_add_(0, mask_ids, h_rows * mask_weights) + h_patch = weighted_mean(h_mask.unflatten(0, (n_global, batch_size)).mean(0)) + + cross = dino.new_zeros(()) + cross_metrics = {"cross": cross, "cross_kl": cross, "cross_pairs": cross} + if self.cross_view_weight > 0 and "cross_a" in batch: + cross, cross_metrics = self._cross_view( + batch, t_patch, global_out["x_norm_patchtokens"], teacher_temperature, + grid, n_tokens, batch_size, weighted_mean, + ) + total = ( self.weights.dino * dino + self.weights.ibot * ibot + self.weights.koleo * koleo + self.weights.gram * gram + + self.cross_view_weight * cross ) return total, { "loss": total.detach(), + "dino_target_entropy": cls_stats["entropy"], + "ibot_target_entropy": patch_stats["entropy"], + "dino_target_top1": cls_stats["top1"], + "ibot_target_top1": patch_stats["top1"], + "dino_prototype_usage": cls_stats["usage"], + "ibot_prototype_usage": patch_stats["usage"], + "dino_kl": (weighted_mean(dino_global) - h_cls).detach(), + "ibot_kl": (ibot - h_patch).detach(), + **{k: v.detach() for k, v in cross_metrics.items()}, "dino": dino.detach(), "dino_global": weighted_mean(dino_global).detach(), "dino_local": weighted_mean(dino_local).detach(), diff --git a/dino-pretraining/mr_dino/probe.py b/dino-pretraining/mr_dino/probe.py new file mode 100644 index 0000000..9108072 --- /dev/null +++ b/dino-pretraining/mr_dino/probe.py @@ -0,0 +1,513 @@ +"""Validation MIL probe for MR DINO checkpoints (the FORA CT-DINO checkpoint protocol). + +A checkpoint is scored without touching train or test data: + +1. ``export``: copy the EMA-teacher backbone out of a DCP training checkpoint into + one small ``.pt`` file, so later steps do not depend on checkpoint rotation. +2. ``extract``: run the frozen backbone over every sequence of every labelled + study of one split (default ``val``). + - Each volume goes through the training preprocessing (same space, spacing and + shape as the checkpoint), cut into training-size tiles (default: the phase-1 + global crop, 64x192x192 voxels). + - Output: the last block's normalised patch tokens. + - Tokens are average-pooled ``--pool`` cells (default 2x2x2) over head-foreground + tokens. Air-only cells are dropped. + - Each worker writes one raw fp16 token file plus one index, so the inode use is + tiny. +3. ``cv``: patient-level K-fold cross-validation of NeuroVFM's ClassifyThenAggregate + head (the MR-RATE MIL head) on the frozen bags. One bag holds all sequences of a + study. Only out-of-fold predictions are kept. +4. ``summarize``: per-class and macro AUROC / AUPRC over all studies. +""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import importlib.util +import json +import math +import os +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .data import ( + ALIGNED_SPACES, + _load_split_ids, + _raw_volume, + discover_raw_aligned, +) + +TEACHER_PREFIX = "model.teacher.backbone." +EXPORT_FORMAT = "mrdino_teacher_backbone_v1" +INDEX_FORMAT = "mrdino_probe_tokens_v1" +REPO_ROOT = Path(__file__).resolve().parents[2] +DEFAULT_LABELS = REPO_ROOT / "contrastive-pretraining/scripts/eval_labels/splits_merged_majority/mrrate_merged_labels.csv" +DEFAULT_SPLITS = REPO_ROOT / "contrastive-pretraining/scripts/eval_labels/splits_merged_majority/splits.csv" + + +# --------------------------------------------------------------------------- checkpoint + +def read_checkpoint_metadata(checkpoint: Path) -> dict: + """Training args/step of a DCP checkpoint directory or an exported teacher file.""" + checkpoint = Path(checkpoint) + if checkpoint.is_dir(): + meta = json.loads((checkpoint / "metadata.json").read_text()) + return {"step": int(meta["step"]), "stage": meta["stage"], "args": meta["args"], "source": str(checkpoint)} + state = torch.load(checkpoint, map_location="cpu", weights_only=False) + if state.get("format") != EXPORT_FORMAT: + raise ValueError(f"{checkpoint} is neither a DCP checkpoint directory nor a {EXPORT_FORMAT} file") + return {k: state[k] for k in ("step", "stage", "args", "source")} + + +def load_teacher_backbone(checkpoint: Path, device: torch.device | str = "cpu"): + """Rebuild the EMA-teacher backbone from a DCP checkpoint (no process group needed) or an export.""" + from .train_ddp import build_backbone + + checkpoint = Path(checkpoint) + meta = read_checkpoint_metadata(checkpoint) + backbone = build_backbone(meta["args"]["arch"]) + if checkpoint.is_dir(): + if not (checkpoint / "COMPLETE").exists(): + raise RuntimeError(f"Incomplete checkpoint: {checkpoint}") + import torch.distributed.checkpoint as dcp + + wanted = {TEACHER_PREFIX + k: torch.empty_like(v) for k, v in backbone.state_dict().items()} + dcp.load(wanted, storage_reader=dcp.FileSystemReader(str(checkpoint)), no_dist=True) + state = {k[len(TEACHER_PREFIX):]: v for k, v in wanted.items()} + else: + state = torch.load(checkpoint, map_location="cpu", weights_only=False)["state_dict"] + backbone.load_state_dict(state, strict=True) + return backbone.to(device).eval().requires_grad_(False), meta + + +def export_teacher(checkpoint: Path, out: Path) -> Path: + backbone, meta = load_teacher_backbone(checkpoint) + out.parent.mkdir(parents=True, exist_ok=True) + tmp = out.with_suffix(out.suffix + ".partial") + torch.save({"format": EXPORT_FORMAT, **meta, "state_dict": backbone.state_dict()}, tmp) + os.replace(tmp, out) + return out + + +# --------------------------------------------------------------------------- tokens + +def _pad_to_multiple(x: torch.Tensor, multiple: tuple[int, int, int], value: float) -> torch.Tensor: + pads = [] + for n, m in reversed(list(zip(x.shape[-3:], multiple))): + pads.extend((0, (-n) % m)) + return F.pad(x, pads, value=value) if any(pads) else x + + +@torch.inference_mode() +def embed_sequence( + backbone, + volume: torch.Tensor, + tile: tuple[int, int, int], + pool: tuple[int, int, int], + patch: tuple[int, int, int] = (2, 16, 16), + fg_margin: float = 0.05, + tiles_per_batch: int = 8, + device: torch.device | str = "cpu", +) -> tuple[torch.Tensor, torch.Tensor]: + """Return pooled foreground tokens ``[M, C]`` (fp16) and their pooled-grid coords ``[M, 3]``. + + Background is the volume minimum (air after MR-RATE's z-score; padding uses it too). + A token counts as foreground when some voxel of its patch is ``fg_margin`` above it. + """ + if any(t % p for t, p in zip(tile, patch)): + raise ValueError(f"Tile {tile} must be a multiple of the patch {patch}") + device = torch.device(device) + background = float(volume.float().min()) + vol = _pad_to_multiple(volume.float(), tile, background) + grid = tuple(n // p for n, p in zip(vol.shape, patch)) + tile_grid = tuple(t // p for t, p in zip(tile, patch)) + starts = [ + (z, y, x) + for z in range(0, vol.shape[0], tile[0]) + for y in range(0, vol.shape[1], tile[1]) + for x in range(0, vol.shape[2], tile[2]) + ] + feats = torch.zeros(*grid, backbone.embed_dim, device=device, dtype=torch.float32) + fg = torch.zeros(grid, device=device, dtype=torch.bool) + autocast = torch.autocast("cuda", dtype=torch.bfloat16) if device.type == "cuda" else torch.autocast("cpu", enabled=False) + for i in range(0, len(starts), tiles_per_batch): + chunk = starts[i : i + tiles_per_batch] + x = torch.stack([vol[z : z + tile[0], y : y + tile[1], xx : xx + tile[2]] for z, y, xx in chunk]) + x = x.to(device)[:, None] + tile_fg = F.max_pool3d(x - background, kernel_size=patch, stride=patch)[:, 0] > fg_margin + keep = tile_fg.flatten(1).any(1) + if not bool(keep.any()): + continue + with autocast: + out = backbone.forward_features(x[keep].to(torch.bfloat16 if device.type == "cuda" else torch.float32)) + tokens = out["x_norm_patchtokens"].float().reshape(-1, *tile_grid, backbone.embed_dim) + for (z, y, xx), t, m in zip([s for s, k in zip(chunk, keep.tolist()) if k], tokens, tile_fg[keep]): + gz, gy, gx = z // patch[0], y // patch[1], xx // patch[2] + feats[gz : gz + tile_grid[0], gy : gy + tile_grid[1], gx : gx + tile_grid[2]] = t + fg[gz : gz + tile_grid[0], gy : gy + tile_grid[1], gx : gx + tile_grid[2]] = m + # foreground-weighted average pooling on the token grid + weight = _pad_to_multiple(fg.float()[None], pool, 0.0)[0] + feats = _pad_to_multiple(feats.permute(3, 0, 1, 2), pool, 0.0) + count = F.avg_pool3d(weight[None, None], pool, pool)[0, 0] + summed = F.avg_pool3d((feats * weight)[None], pool, pool)[0] + cells = count > 0 + pooled = (summed[:, cells] / count[cells]).T + coords = cells.nonzero().to(torch.int16) + return pooled.to(torch.float16).cpu(), coords.cpu() + + +def worker_identity() -> tuple[int, int]: + rank = int(os.environ.get("RANK", os.environ.get("SLURM_PROCID", 0))) + world = int(os.environ.get("WORLD_SIZE", os.environ.get("SLURM_NTASKS", 1))) + return rank, world + + +def read_labels(labels_csv: Path) -> tuple[list[str], dict[str, np.ndarray]]: + with open(labels_csv, newline="") as handle: + reader = csv.DictReader(handle) + names = [c for c in reader.fieldnames if c not in {"study_uid", "subject_id"}] + key = "study_uid" if "study_uid" in reader.fieldnames else "subject_id" + labels = {row[key].strip(): np.array([float(row[c]) for c in names], dtype=np.float32) for row in reader} + return names, labels + + +def probe_studies(args) -> list[dict]: + """Labelled studies of the split, sorted by uid (identical on every worker).""" + _, labels = read_labels(Path(args.labels_csv)) + selected = _load_split_ids(str(args.splits_csv), args.split) + selected = {uid for uid in selected if uid in labels} + studies = sorted(discover_raw_aligned(args.data_folder, selected, args.space), key=lambda s: s["study_uid"]) + if args.max_studies: + studies = studies[: args.max_studies] + if not studies: + raise RuntimeError(f"No labelled {args.split} studies under {args.data_folder}") + return studies + + +class _StudyVolumes(torch.utils.data.Dataset): + def __init__(self, studies, target_shape, target_spacing, posterior_shift_mm): + self.studies, self.shape, self.spacing, self.shift = studies, target_shape, target_spacing, posterior_shift_mm + + def __len__(self): + return len(self.studies) + + def __getitem__(self, i): + study = self.studies[i] + volumes = [_raw_volume(p, self.shape, self.spacing, self.shift) for p in study["image_paths"]] + return study["study_uid"], [Path(p).name for p in study["image_paths"]], volumes + + +def extract(args) -> Path: + rank, world = worker_identity() + device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu")) + backbone, meta = load_teacher_backbone(Path(args.checkpoint), device) + train = meta["args"] + args.space = args.space or train.get("space", "atlas_space") # pre-coreg checkpoints were atlas + if args.space != train.get("space", "atlas_space") and not args.allow_space_mismatch: + raise ValueError(f"Checkpoint was trained on {train.get('space')!r}, probe asked for {args.space!r}") + target_shape = tuple(train["target_shape"]) + target_spacing = tuple(train["target_spacing"]) + out = Path(args.features_dir) + out.mkdir(parents=True, exist_ok=True) + index_path = out / f"index_rank{rank:04d}.npz" + if index_path.exists(): + print(f"[probe] rank {rank}: {index_path} exists, skipping", flush=True) + return index_path + studies = probe_studies(args)[rank::world] + loader = torch.utils.data.DataLoader( + _StudyVolumes(studies, target_shape, target_spacing, float(train["posterior_shift_mm"])), + batch_size=None, shuffle=False, num_workers=args.workers, + prefetch_factor=1 if args.workers else None, persistent_workers=False, + ) + token_path = out / f"tokens_rank{rank:04d}.f16" + uids, starts, counts, seq_ids, coords, names = [], [], [], [], [], [] + total, t0 = 0, time.time() + with open(token_path, "wb") as sink: + for i, (uid, seq_names, volumes) in enumerate(loader): + start = total + for s, volume in enumerate(volumes): + tokens, xyz = embed_sequence( + backbone, volume, tuple(args.tile), tuple(args.pool), + fg_margin=args.fg_margin, tiles_per_batch=args.tiles_per_batch, device=device, + ) + sink.write(tokens.numpy().tobytes()) + total += len(tokens) + seq_ids.append(np.full(len(tokens), s, dtype=np.int16)) + coords.append(xyz.numpy()) + if total == start: + raise RuntimeError(f"Study {uid} produced no foreground tokens") + uids.append(uid) + starts.append(start) + counts.append(total - start) + names.append(json.dumps(seq_names)) + if rank == 0 and (i % 20 == 0 or i + 1 == len(studies)): + rate = (i + 1) / max(1e-6, time.time() - t0) + print(f"[probe] rank 0: {i + 1}/{len(studies)} studies, {total:,} tokens, {rate:.2f} studies/s", flush=True) + tmp = out / f"index_rank{rank:04d}.partial.npz" + np.savez( + tmp, + format=INDEX_FORMAT, dim=backbone.embed_dim, token_file=token_path.name, + study_uid=np.array(uids), start=np.array(starts, dtype=np.int64), count=np.array(counts, dtype=np.int64), + seq_id=np.concatenate(seq_ids) if seq_ids else np.zeros(0, np.int16), + coords=np.concatenate(coords) if coords else np.zeros((0, 3), np.int16), + sequence_names=np.array(names), + checkpoint_step=meta["step"], checkpoint_stage=meta["stage"], checkpoint_source=meta["source"], + space=args.space, tile=np.array(args.tile), pool=np.array(args.pool), + ) + os.replace(tmp, index_path) + return index_path + + +# --------------------------------------------------------------------------- MIL head + +def mil_head_class(): + """NeuroVFM ClassifyThenAggregate from MR-RATE's MIL probe (single source of truth).""" + path = Path(os.environ.get("MRRATE_MIL_PROBE_MODULE", REPO_ROOT / "contrastive-pretraining/scripts/mil_probe.py")) + spec = importlib.util.spec_from_file_location("mrrate_mil_probe", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.ClassifyThenAggregate + + +class TokenBags: + """All workers' token files as memory maps, one bag per study.""" + + def __init__(self, features_dir: Path): + self.uids, self.bags, self.dim = [], [], None + indexes = sorted(Path(features_dir).glob("index_rank*.npz")) + if not indexes: + raise FileNotFoundError(f"No index_rank*.npz under {features_dir}") + self.meta = {} + for path in indexes: + with np.load(path, allow_pickle=False) as index: + if str(index["format"]) != INDEX_FORMAT: + raise ValueError(f"{path} is not a {INDEX_FORMAT} index") + dim = int(index["dim"]) + if self.dim not in (None, dim): + raise ValueError("Workers disagree on the token width") + self.dim = dim + tokens = np.memmap(path.parent / str(index["token_file"]), dtype=np.float16, mode="r").reshape(-1, dim) + for uid, start, count in zip(index["study_uid"], index["start"], index["count"]): + self.uids.append(str(uid)) + self.bags.append((tokens, int(start), int(count))) + self.meta = {k: index[k].tolist() for k in ("checkpoint_step", "checkpoint_stage", "space", "tile", "pool")} + if len(set(self.uids)) != len(self.uids): + raise ValueError("A study appears in more than one worker shard") + + def __len__(self): + return len(self.uids) + + def tokens(self, i: int) -> np.ndarray: + array, start, count = self.bags[i] + return np.asarray(array[start : start + count]) + + +def patient_folds(uids: list[str], patient_of: dict[str, str], k: int, seed: int) -> np.ndarray: + """Deterministic patient-level fold id per study (all studies of a patient share a fold).""" + patients = sorted({patient_of.get(uid, uid) for uid in uids}) + order = sorted(patients, key=lambda p: hashlib.sha256(f"{seed}:{p}".encode()).hexdigest()) + fold_of = {p: i % k for i, p in enumerate(order)} + return np.array([fold_of[patient_of.get(uid, uid)] for uid in uids]) + + +def read_patients(splits_csv: Path) -> dict[str, str]: + with open(splits_csv, newline="") as handle: + return {row["study_uid"].strip(): row.get("patient_uid", row["study_uid"]).strip() for row in csv.DictReader(handle)} + + +def _batch(bags: TokenBags, items, device): + arrays = [bags.tokens(i) for i in items] + lengths = torch.tensor([0] + [len(a) for a in arrays]) + tokens = torch.from_numpy(np.concatenate(arrays)).to(device, non_blocking=True) + return tokens, torch.cumsum(lengths, 0).to(device) + + +def run_fold(args) -> Path: + device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu")) + bags = TokenBags(Path(args.features_dir)) + names, labels = read_labels(Path(args.labels_csv)) + y = np.stack([labels[uid] for uid in bags.uids]) + folds = patient_folds(bags.uids, read_patients(Path(args.splits_csv)), args.folds, args.seed) + train = np.flatnonzero(folds != args.fold) + held = np.flatnonzero(folds == args.fold) + torch.manual_seed(args.seed + args.fold) + rng = np.random.default_rng(args.seed + args.fold) + + head = mil_head_class()(bags.dim, len(names), hidden_dim=args.hidden_dim, mlp_hidden_dims=(args.mlp_hidden_dim,)).to(device) + positives = y[train].sum(0) + pos_weight = np.clip((len(train) - positives) / np.clip(positives, 1.0, None), 1.0, 100.0) + loss_fn = nn.BCEWithLogitsLoss(pos_weight=torch.tensor(pos_weight, dtype=torch.float32, device=device)) + optimizer = torch.optim.AdamW(head.parameters(), lr=args.lr, weight_decay=args.weight_decay) + steps = args.epochs * math.ceil(len(train) / args.batch_size) + warmup = max(1, int(0.1 * steps)) + schedule = torch.optim.lr_scheduler.LambdaLR( + optimizer, lambda s: (s + 1) / warmup if s < warmup else 0.5 * (1 + math.cos(math.pi * (s - warmup) / max(1, steps - warmup))) + ) + autocast = torch.autocast("cuda", dtype=torch.bfloat16) if device.type == "cuda" else torch.autocast("cpu", enabled=False) + head.train() + for epoch in range(args.epochs): + order = rng.permutation(train) + for b in range(0, len(order), args.batch_size): + items = order[b : b + args.batch_size] + tokens, cu = _batch(bags, items, device) + with autocast: + logits = head(tokens.float(), cu) + loss = loss_fn(logits.float(), torch.from_numpy(y[items]).to(device)) + optimizer.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(head.parameters(), 1.0) + optimizer.step() + schedule.step() + print(f"[probe] fold {args.fold} epoch {epoch + 1}/{args.epochs} loss {float(loss.detach()):.4f}", flush=True) + head.eval() + out_logits = [] + with torch.inference_mode(): + for b in range(0, len(held), args.batch_size): + items = held[b : b + args.batch_size] + tokens, cu = _batch(bags, items, device) + with autocast: + out_logits.append(head(tokens.float(), cu).float().cpu().numpy()) + out = Path(args.out_dir) + out.mkdir(parents=True, exist_ok=True) + path = out / f"oof_fold{args.fold}.npz" + np.savez(path, study_uid=np.array([bags.uids[i] for i in held]), logits=np.concatenate(out_logits), labels=y[held]) + return path + + +# --------------------------------------------------------------------------- metrics + +def auroc(target: np.ndarray, score: np.ndarray) -> float | None: + """Mann-Whitney AUROC with average ranks for ties.""" + positive = target == 1 + n1, n0 = int(positive.sum()), int((~positive).sum()) + if n1 == 0 or n0 == 0: + return None + _, inverse, counts = np.unique(score, return_inverse=True, return_counts=True) + average_rank = np.cumsum(counts) - (counts - 1) / 2.0 + ranks = average_rank[inverse] + return float((ranks[positive].sum() - n1 * (n1 + 1) / 2.0) / (n1 * n0)) + + +def average_precision(target: np.ndarray, score: np.ndarray) -> float | None: + if not (target == 1).any(): + return None + hits = target[np.argsort(-score, kind="mergesort")] == 1 + precision = np.cumsum(hits) / np.arange(1, len(hits) + 1) + return float(precision[hits].mean()) + + +def summarize(args) -> dict: + out = Path(args.out_dir) + names, _ = read_labels(Path(args.labels_csv)) + parts = [np.load(out / f"oof_fold{f}.npz") for f in range(args.folds)] + uids = np.concatenate([p["study_uid"] for p in parts]) + if len(set(uids.tolist())) != len(uids): + raise ValueError("A study was predicted by more than one fold") + logits = np.concatenate([p["logits"] for p in parts]) + labels = np.concatenate([p["labels"] for p in parts]) + rows = [] + for c, name in enumerate(names): + rows.append({ + "label": name, "positives": int(labels[:, c].sum()), "studies": int(len(labels)), + "auroc": auroc(labels[:, c], logits[:, c]), "auprc": average_precision(labels[:, c], logits[:, c]), + }) + valid = [r for r in rows if r["auroc"] is not None] + meta = TokenBags(Path(args.features_dir)).meta if args.features_dir and Path(args.features_dir).exists() else {} + fold_auroc = [] + for p in parts: + per = [auroc(p["labels"][:, c], p["logits"][:, c]) for c in range(len(names))] + per = [v for v in per if v is not None] + fold_auroc.append(float(np.mean(per)) if per else None) + result = { + "macro_auroc": float(np.mean([r["auroc"] for r in valid])), + "macro_auprc": float(np.mean([r["auprc"] for r in valid])), + "fold_macro_auroc": fold_auroc, + "studies": int(len(labels)), + "folds": args.folds, + "labels_csv": str(args.labels_csv), + "per_class": rows, + **meta, + } + (out / "results.json").write_text(json.dumps(result, indent=2)) + with open(out / "per_class.csv", "w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + print( + f"[probe] step {meta.get('checkpoint_step', '?')}: macro AUROC {result['macro_auroc']:.4f} " + f"AUPRC {result['macro_auprc']:.4f} over {result['studies']} studies " + f"(folds {', '.join(f'{v:.3f}' for v in fold_auroc if v is not None)})", + flush=True, + ) + return result + + +# --------------------------------------------------------------------------- CLI + +def parse_args(argv=None) -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = p.add_subparsers(dest="command", required=True) + + e = sub.add_parser("export", help="Copy the EMA-teacher backbone out of a DCP checkpoint") + e.add_argument("--checkpoint", required=True) + e.add_argument("--out", required=True) + + common = argparse.ArgumentParser(add_help=False) + common.add_argument("--labels-csv", default=str(DEFAULT_LABELS), help="study_uid + binary label columns (default: 14 merged groups)") + common.add_argument("--splits-csv", default=str(DEFAULT_SPLITS), help="study_uid, patient_uid, split") + common.add_argument("--features-dir", required=True) + common.add_argument("--device", default=None) + + x = sub.add_parser("extract", parents=[common], help="Frozen tokens of every labelled study of a split") + x.add_argument("--checkpoint", required=True, help="DCP checkpoint directory or exported teacher .pt") + x.add_argument("--data-folder", required=True, help="Extracted MR-RATE tree: batchXX//coreg_img (or atlas_img)") + x.add_argument("--space", default=None, choices=ALIGNED_SPACES, help="Default: the checkpoint's training space") + x.add_argument("--allow-space-mismatch", action="store_true") + x.add_argument("--split", default="val") + x.add_argument("--tile", type=int, nargs=3, default=(64, 192, 192), help="Voxels; default = phase-1 global crop") + x.add_argument("--pool", type=int, nargs=3, default=(2, 2, 2), help="Token-grid pooling cell") + x.add_argument("--fg-margin", type=float, default=0.05, help="Foreground: patch max above the volume minimum") + x.add_argument("--tiles-per-batch", type=int, default=8) + x.add_argument("--workers", type=int, default=6) + x.add_argument("--max-studies", type=int, default=None, help="Smoke only") + + c = sub.add_parser("cv", parents=[common], help="Train/predict one patient-level fold") + c.add_argument("--out-dir", required=True) + c.add_argument("--fold", type=int, required=True) + c.add_argument("--folds", type=int, default=5) + c.add_argument("--seed", type=int, default=42) + c.add_argument("--epochs", type=int, default=8) + c.add_argument("--batch-size", type=int, default=8) + c.add_argument("--lr", type=float, default=2e-3) + c.add_argument("--weight-decay", type=float, default=0.05) + c.add_argument("--hidden-dim", type=int, default=512) + c.add_argument("--mlp-hidden-dim", type=int, default=384) + + s = sub.add_parser("summarize", parents=[common], help="Out-of-fold macro/per-class AUROC and AUPRC") + s.add_argument("--out-dir", required=True) + s.add_argument("--folds", type=int, default=5) + return p.parse_args(argv) + + +def main(argv=None) -> int: + args = parse_args(argv) + if args.command == "export": + print(export_teacher(Path(args.checkpoint), Path(args.out))) + elif args.command == "extract": + print(extract(args)) + elif args.command == "cv": + print(run_fold(args)) + else: + summarize(args) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dino-pretraining/mr_dino/recipe.py b/dino-pretraining/mr_dino/recipe.py new file mode 100644 index 0000000..08c65d2 --- /dev/null +++ b/dino-pretraining/mr_dino/recipe.py @@ -0,0 +1,282 @@ +"""Recipe fixes ported from the FORA CT DINOv3 pipeline (v13/v14, Sep 2026). + +The CT pipeline started from the same DINOv3-7B recipe as this module and failed in ways we +diagnosed and fixed one by one; every piece below is what the working CT ViT-L run uses: + +* cosine prototypes on both heads (``CosineLinear``): DINOv3's plain linear prototype layer + (std 0.02) never sharpened at small learning rates, so the Sinkhorn targets stayed flat + (iBOT target entropy 8.0/11.5 nats, top-1 0.006) for thousands of steps; +* position-binned Sinkhorn for the patch head (CAPI, arXiv 2502.08769): the patch targets are + balanced within spatial bins of the crop instead of across all masked tokens at once; +* DINOv2 optimizer groups: layer-wise lr decay, patch-embed lr x0.2, no weight decay on biases, + 1-D parameters and tokens, and the prototype layers frozen for the first steps; +* cross-crop twin tokens: with the two global crops aligned on the patch lattice every token in + their overlap has an exact twin, which the objective uses to stop the head-facing block from + encoding the token's position inside the crop (measured in CT: last-block position leakage + 0.201 -> 0.094, crop-quadrant prototype fields gone). +""" + +from __future__ import annotations + +import json +import math +import os + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch import Tensor, nn + +TOKEN_NAMES = ("cls_token", "mask_token", "storage_tokens") + + +# ----------------------------------------------------------------------------- prototypes +class CosineLinear(nn.Linear): + """Bias-free prototype layer computing cosine logits (weights L2-normalized per prototype). + + DINOv3's ``DINOHead`` already L2-normalizes the bottleneck features before ``last_layer``, + so the logits are true cosine similarities. The parameter name stays ``last_layer.weight``, + keeping checkpoints of the plain layer loadable.""" + + def forward(self, x: Tensor) -> Tensor: + return F.linear(x, F.normalize(self.weight, dim=1, eps=1e-12)) + + +def normalize_prototype_layer(head: nn.Module) -> None: + layer = head.last_layer + if not isinstance(layer, nn.Linear) or layer.bias is not None: + raise ValueError("Expected a bias-free linear prototype layer") + if isinstance(layer, CosineLinear): + return + cosine = CosineLinear(layer.in_features, layer.out_features, bias=False, + device=layer.weight.device, dtype=layer.weight.dtype) + with torch.no_grad(): + cosine.weight.copy_(layer.weight) + head.last_layer = cosine + + +# ----------------------------------------------------------------------------- Sinkhorn bins +def position_bin_ids(token_indices: Tensor, n_tokens: int, grid, bins) -> Tensor: + """Spatial bin id in [0, prod(bins)) of flat token indices (modulo the crop's token count).""" + gz, gy, gx = (int(v) for v in grid) + if gz * gy * gx != int(n_tokens): + raise ValueError(f"position grid {grid} does not match {n_tokens} tokens per crop") + bz, by, bx = (int(v) for v in bins) + if not (1 <= bz <= gz and 1 <= by <= gy and 1 <= bx <= gx): + raise ValueError(f"position bins {bins} must fit the token grid {grid}") + pos = token_indices % int(n_tokens) + z = pos // (gy * gx) + y = (pos // gx) % gy + x = pos % gx + return ((z * bz) // gz) * (by * bx) + ((y * by) // gy) * bx + (x * bx) // gx + + +def binned_sinkhorn(base, logits: Tensor, temperature: float, bin_ids: Tensor, n_bins: int, + iterations: int = 3) -> Tensor: + """Run the distributed Sinkhorn ``base`` once per spatial bin; every rank visits every + non-empty bin in the same order (bin counts are all-reduced), so collectives stay matched.""" + if bin_ids.shape[0] != logits.shape[0]: + raise ValueError("bin ids must align with the logit rows") + out = torch.empty(logits.shape, dtype=torch.float32, device=logits.device) + counts = torch.bincount(bin_ids, minlength=n_bins).to(torch.int64).to(logits.device) + if dist.is_initialized(): + dist.all_reduce(counts, op=dist.ReduceOp.SUM) + for b in range(n_bins): + if int(counts[b]) == 0: + continue + rows = (bin_ids == b).nonzero().flatten() + q = base(logits.index_select(0, rows), temperature, iterations) + out.index_copy_(0, rows, q) + return out + + +# ----------------------------------------------------------------------------- optimizer +def _block_index(name: str): + parts = name.split(".") + for i, part in enumerate(parts[:-1]): + if part == "blocks" and parts[i + 1].isdigit(): + return int(parts[i + 1]) + return None + + +def build_param_groups(student: nn.Module, *, layerwise_decay: float, patch_embed_lr_mult: float, + n_blocks: int) -> list[dict]: + """DINOv2 ``get_params_groups_with_decay`` for the 3-D student (backbone + two heads). + + lr_mult: backbone blocks[i] get decay^(n_blocks - i); patch embedding, tokens and RoPE get + decay^(n_blocks + 1) (times ``patch_embed_lr_mult`` for the patch embedding); final norms + get 1; heads get 1. wd_mult: 0 for biases, 1-D parameters and tokens. Parameters that share + (lr_mult, wd_mult) are merged into one group.""" + if not 0.0 < layerwise_decay <= 1.0 or patch_embed_lr_mult <= 0 or n_blocks < 1: + raise ValueError("Require 0 < layerwise_decay <= 1, patch_embed_lr_mult > 0, n_blocks >= 1") + merged: dict[tuple[float, float], dict] = {} + for name, param in student.named_parameters(): + if not param.requires_grad: + continue + clean = name.replace("_fsdp_wrapped_module.", "").replace("_checkpoint_wrapped_module.", "").replace("_orig_mod.", "") + lr_mult, wd_mult = 1.0, 1.0 + if clean.startswith("backbone."): + block = _block_index(clean) + if block is not None: + layer_id = block + 1 + elif "patch_embed" in clean or any(t in clean for t in TOKEN_NAMES) or "rope" in clean: + layer_id = 0 + else: + layer_id = n_blocks + 1 + lr_mult = float(layerwise_decay) ** (n_blocks + 1 - layer_id) + if "patch_embed" in clean: + lr_mult *= float(patch_embed_lr_mult) + if param.ndim <= 1 or clean.endswith(".bias") or any(t in clean for t in TOKEN_NAMES): + wd_mult = 0.0 + key = (round(lr_mult, 12), wd_mult) + group = merged.setdefault(key, dict(params=[], lr_mult=float(lr_mult), wd_mult=float(wd_mult), names=[])) + group["params"].append(param) + group["names"].append(clean) + groups = [merged[k] for k in sorted(merged, key=lambda k: (-k[0], -k[1]))] + if os.environ.get("RANK", "0") == "0": + print(json.dumps(dict(event="param_groups", groups=[ + dict(lr_mult=g["lr_mult"], wd_mult=g["wd_mult"], n=len(g["params"]), example=g["names"][0]) + for g in groups])), flush=True) + return groups + + +class GroupedAdamW(torch.optim.AdamW): + """AdamW applying each group's ``lr_mult``/``wd_mult`` transiently inside ``step`` so the + trainer keeps writing the plain scheduled lr / weight decay into every group.""" + + @torch.no_grad() + def step(self, closure=None): + saved = [(g["lr"], g["weight_decay"]) for g in self.param_groups] + for g in self.param_groups: + g["lr"] = g["lr"] * float(g.get("lr_mult", 1.0)) + g["weight_decay"] = g["weight_decay"] * float(g.get("wd_mult", 1.0)) + try: + return super().step(closure) + finally: + for g, (lr, wd) in zip(self.param_groups, saved): + g["lr"], g["weight_decay"] = lr, wd + + +def make_optimizer(student: nn.Module, *, lr: float, weight_decay: float, betas, layerwise_decay: float, + patch_embed_lr_mult: float, n_blocks: int, fused: bool = True) -> GroupedAdamW: + groups = build_param_groups(student, layerwise_decay=layerwise_decay, + patch_embed_lr_mult=patch_embed_lr_mult, n_blocks=n_blocks) + for g in groups: + g.pop("names") + return GroupedAdamW(groups, lr=lr, betas=tuple(betas), weight_decay=weight_decay, fused=fused) + + +def cancel_last_layer_gradients(student: nn.Module) -> int: + """DINO/DINOv2 ``freeze_last_layer``: drop the prototype-layer gradients of both heads.""" + count = 0 + for name, param in student.named_parameters(): + if "last_layer" in name and param.grad is not None: + param.grad = None + count += 1 + return count + + +def missing_frozen_optimizer_keys(ckpt_dir, model: nn.Module, optimizer: torch.optim.Optimizer) -> list[str]: + """Checkpoint keys the resume would miss that are legitimately absent: optimizer state of the ``last_layer`` + prototype weights, which never receive a gradient during ``freeze_last_layer`` (``cancel_last_layer_gradients`` + sets ``grad=None``, so AdamW never creates their state and checkpoints saved in the first + ``--freeze-last-layer-steps`` steps do not contain it). Any OTHER missing model/optimizer key raises, so a + partial load can never silently skip real weights. Returns the allowed missing keys (empty = strict load).""" + from torch.distributed.checkpoint import FileSystemReader + from torch.distributed.checkpoint._nested_dict import flatten_state_dict + from torch.distributed.checkpoint.state_dict import get_model_state_dict, get_optimizer_state_dict + saved = set(FileSystemReader(str(ckpt_dir)).read_metadata().state_dict_metadata) + want, _ = flatten_state_dict({"model": get_model_state_dict(model), + "optimizer": get_optimizer_state_dict(model, optimizer)}) + missing = [k for k in want if k not in saved] + bad = [k for k in missing if not (k.startswith("optimizer.state.") and ".last_layer." in k)] + if bad: + raise RuntimeError(f"Checkpoint {ckpt_dir} lacks {len(bad)} required keys, e.g. {bad[:5]}") + return missing + + +# ----------------------------------------------------------------------------- twin tokens +def snap_second_start(origin, second, container_shape, crop_shape, patch): + """Move ``second`` so that (second - origin) is a multiple of the patch on every axis, + staying inside the container.""" + out = [] + for o, s, have, want, p in zip(origin, second, container_shape, crop_shape, patch): + k = int(round((s - o) / p)) + lo = -(o // p) + hi = (max(have - want, 0) - o) // p + if hi < lo: + hi = lo + out.append(o + min(max(k, lo), hi) * p) + return tuple(out) + + +def token_offset(g0, g1, patch): + d = [] + for a, b, p in zip(g0, g1, patch): + if (b - a) % p: + raise ValueError("global crops are not aligned on the patch lattice") + d.append((b - a) // p) + return tuple(d) + + +def overlap_pairs(offset, grid) -> tuple[Tensor, Tensor]: + """Flat token indices (crop A, crop B) of every token covering the same voxels in both crops.""" + ranges = [] + for d, g in zip(offset, grid): + lo, hi = max(0, d), min(g, g + d) + if hi <= lo: + return torch.empty(0, dtype=torch.long), torch.empty(0, dtype=torch.long) + ranges.append(torch.arange(lo, hi)) + z, y, x = torch.meshgrid(*ranges, indexing="ij") + gz, gy, gx = grid + a = (z * gy + y) * gx + x + b = ((z - offset[0]) * gy + (y - offset[1])) * gx + (x - offset[2]) + return a.flatten(), b.flatten() + + +def padded_head(head: nn.Module, tokens: Tensor, chunk: int = 4096) -> tuple[Tensor, Tensor]: + """Run an FSDP-sharded head on a per-rank-variable token count with rank-identical call shapes. + Returns (logits[:local], anchor); anchor keeps a zero graph edge on ranks without tokens.""" + local = int(tokens.shape[0]) + maximum = torch.tensor(local, device=tokens.device) + if dist.is_initialized(): + dist.all_reduce(maximum, op=dist.ReduceOp.MAX) + maximum = int(maximum.item()) + parts, out = [], None + anchor = tokens.sum() * 0.0 + for start in range(0, maximum, chunk): + stop = min(start + chunk, maximum) + size = stop - start + valid = max(0, min(stop, local) - start) + x = tokens[start:start + valid] + if valid < size: + x = F.pad(x, (0, 0, 0, size - valid)) + out = head(x) + if valid == 0: + anchor = anchor + out.sum() * 0.0 + continue + parts.append(out[:valid]) + if parts: + logits = torch.cat(parts) + elif out is not None: + logits = out[:0] + else: + logits = tokens.new_zeros((0, head.last_layer.weight.shape[0])) + return logits, anchor + + +def entropy(q: Tensor) -> Tensor: + return -(q * q.clamp_min(1e-12).log()).sum(-1) + + +def target_stats(q: Tensor) -> dict[str, Tensor]: + """Teacher-target diagnostics: entropy, mean top-1 probability, fraction of prototypes that win.""" + if q.shape[0] == 0: + z = q.new_zeros(()) + return dict(entropy=z, top1=z, usage=z) + return dict(entropy=entropy(q).mean(), top1=q.max(-1).values.mean(), + usage=torch.tensor(q.argmax(-1).unique().numel() / q.shape[1], device=q.device)) + + +def sqrt_scaled_lr(base_lr: float, global_batch: int, reference: int = 1024) -> float: + return base_lr * math.sqrt(global_batch / reference) diff --git a/dino-pretraining/mr_dino/train_ddp.py b/dino-pretraining/mr_dino/train_ddp.py index fa75c74..7f1accb 100644 --- a/dino-pretraining/mr_dino/train_ddp.py +++ b/dino-pretraining/mr_dino/train_ddp.py @@ -1,4 +1,4 @@ -"""Distributed, resumable 3-D DINOv3 training for atlas-registered MR-RATE.""" +"""Distributed, resumable 3-D DINOv3 training for aligned (coreg/atlas) MR-RATE.""" from __future__ import annotations @@ -23,47 +23,71 @@ from torch.utils.data import DataLoader from .data import ( - MRAtlasDINO3DDataset, + ALIGNED_SPACES, + DEFAULT_SPACE, + MRAlignedDINO3DDataset, InfiniteStudySampler, collate_dino3d, - discover_atlas_cache, - discover_raw_atlas_split, + discover_aligned_cache, + discover_raw_aligned_split, npz_volume_shape, stage_crop_spec, ) from .model import DinoVisionTransformer3D, vit_7b_3d, vit_hplus_3d, vit_large_3d +from .recipe import cancel_last_layer_gradients, make_optimizer from .objective import DINO3DLearner, LossWeights +CHECKPOINT_FORMAT = "mrrate_aligned_dinov3d_full_v1" +# Checkpoints written before the coreg switch: same contents, always atlas space. +_LEGACY_ATLAS_FORMATS = { + "mrrate_aligned_dinov3d_full_v1": "mrrate_atlas_dinov3d_full_v1", + "mrrate_aligned_dinov3d_fsdp2_v1": "mrrate_atlas_dinov3d_fsdp2_v1", +} + + +def check_checkpoint_space(fmt: str | None, saved_args: dict, space: str, expected: str) -> None: + """Refuse to resume across MR spaces.""" + if fmt == expected: + saved_space = saved_args.get("space") + elif fmt == _LEGACY_ATLAS_FORMATS.get(expected): + saved_space = "atlas_space" + else: + raise RuntimeError(f"Checkpoint format {fmt!r} is not an MR DINO checkpoint ({expected})") + if saved_space != space: + raise RuntimeError(f"Checkpoint was trained on {saved_space!r}; this run uses {space!r}") + + def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description=__doc__) source = p.add_mutually_exclusive_group(required=True) - source.add_argument("--preprocessed-dir", help="MR-RATE cache root containing atlas_space/*.npz") + source.add_argument("--preprocessed-dir", help="MR-RATE cache root containing /*.npz") source.add_argument( "--data-folder", - help="MR-RATE-atlas NIfTI tree; uses the same live loader as previous MIL training", + help="MR-RATE-coreg (or -atlas) NIfTI tree; uses the same live loader as previous MIL training", ) p.add_argument("--splits-csv") p.add_argument("--split", default="train") - p.add_argument("--space", default="atlas_space", choices=("atlas_space",)) + p.add_argument("--space", default=DEFAULT_SPACE, choices=ALIGNED_SPACES) p.add_argument("--output-dir", required=True) p.add_argument("--stage", choices=("pretrain", "gram", "highres"), default="pretrain") - p.add_argument("--arch", choices=("tiny", "large", "hplus"), default="hplus") - p.add_argument("--steps", type=int, default=100_000) + p.add_argument("--arch", choices=("tiny", "large", "hplus"), default="large") + p.add_argument("--steps", type=int, default=50_000) p.add_argument("--batch-size", type=int, default=1) p.add_argument("--workers", type=int, default=2) p.add_argument("--seed", type=int, default=3407) - p.add_argument("--prototypes", type=int, default=65_536) - p.add_argument("--head-hidden-dim", type=int, default=4096) - p.add_argument("--lr", type=float, default=2e-4, help="Base LR before sqrt(global_batch/1024) scaling") + p.add_argument("--prototypes", type=int, default=16_384) + p.add_argument("--head-hidden-dim", type=int, default=2048) + p.add_argument("--bottleneck-dim", type=int, default=256) + p.add_argument("--lr", type=float, default=1e-3, help="Base LR before sqrt(global_batch/1024) scaling (DINOv2 ViT-L)") p.add_argument("--min-lr", type=float, default=1e-6) - p.add_argument("--warmup-steps", type=int, default=5_000) + p.add_argument("--warmup-steps", type=int, default=2_500) p.add_argument("--weight-decay", type=float, default=0.04) p.add_argument("--weight-decay-end", type=float, default=0.4) - p.add_argument("--teacher-momentum", type=float, default=0.994) + p.add_argument("--teacher-momentum", type=float, default=0.992) p.add_argument("--teacher-temperature", type=float, default=0.07) p.add_argument("--teacher-warmup-temperature", type=float, default=0.04) - p.add_argument("--teacher-warmup-steps", type=int, default=5_000) + p.add_argument("--teacher-warmup-steps", type=int, default=10_000) p.add_argument("--mask-min", type=float, default=0.1) p.add_argument("--mask-max", type=float, default=0.5) p.add_argument("--ibot-weight", type=float, default=1.0) @@ -78,15 +102,25 @@ def parse_args() -> argparse.Namespace: p.add_argument("--target-spacing", type=float, nargs=3, default=(1.0, 0.5, 0.5)) p.add_argument("--target-shape", type=int, nargs=3, default=(256, 384, 384)) p.add_argument("--posterior-shift-mm", type=float, default=15.0) - p.add_argument("--cross-sequence-probability", type=float, default=0.75) + p.add_argument("--cross-sequence-probability", type=float, default=0.25) p.add_argument("--candidate-trials", type=int, default=12) - p.add_argument("--save-every", type=int, default=250) - p.add_argument("--keep-checkpoints", type=int, default=3) + p.add_argument("--save-every", type=int, default=500) + p.add_argument("--keep-checkpoints", type=int, default=6) p.add_argument("--log-every", type=int, default=10) p.add_argument("--resume", default="latest", help="latest, none, or a checkpoint path") p.add_argument("--activation-checkpointing", action=argparse.BooleanOptionalAction, default=True) p.add_argument("--max-studies", type=int, default=None, help="Smoke only: cap studies assigned to each rank") p.add_argument("--deadline-margin-seconds", type=int, default=900) + # FORA recipe fixes (see mr_dino/recipe.py) + p.add_argument("--global-overlap", type=float, default=0.25) + p.add_argument("--normalize-prototypes", action=argparse.BooleanOptionalAction, default=True) + p.add_argument("--position-bins", type=int, nargs=3, default=(2, 2, 2)) + p.add_argument("--cross-view-weight", type=float, default=1.0) + p.add_argument("--cross-view-sequences", choices=("same", "any", "none"), default="same") + p.add_argument("--layerwise-decay", type=float, default=0.9) + p.add_argument("--patch-embed-lr-mult", type=float, default=0.2) + p.add_argument("--freeze-last-layer-steps", type=int, default=1250) + p.add_argument("--clip-grad", type=float, default=3.0) return p.parse_args() @@ -107,9 +141,9 @@ def read_train_files( preprocessed_dir: str, splits_csv: str | None, split: str, - space: str = "atlas_space", + space: str = DEFAULT_SPACE, ) -> list[str]: - return [sample["cache_path"] for sample in discover_atlas_cache( + return [sample["cache_path"] for sample in discover_aligned_cache( preprocessed_dir, splits_csv=splits_csv, split=split, space=space )] @@ -136,6 +170,7 @@ def raw_fingerprint( target_spacing: tuple[float, float, float] = (1.0, 0.5, 0.5), target_shape: tuple[int, int, int] = (256, 384, 384), posterior_shift_mm: float = 15.0, + space: str = DEFAULT_SPACE, ) -> str: """Fingerprint raw studies without tying exact resume to an absolute mount path.""" records = [] @@ -151,7 +186,7 @@ def raw_fingerprint( "world": world, "max_studies_per_rank": max_studies_per_rank, "preprocessing": { - "space": "atlas_space", + "space": space, "normalizer": "zscore", "target_spacing": list(target_spacing), "target_shape": list(target_shape), @@ -216,16 +251,17 @@ def distributed_raw_assignment( target_spacing: tuple[float, float, float], target_shape: tuple[int, int, int], posterior_shift_mm: float, + space: str = DEFAULT_SPACE, ) -> tuple[list[dict], str]: """Discover once on rank 0 and publish one small assignment file per rank.""" assignment_dir = output / "dataset_assignments" metadata_path = assignment_dir / "metadata.json" if rank == 0: - samples = discover_raw_atlas_split( - data_folder, splits_csv=splits_csv, split=split + samples = discover_raw_aligned_split( + data_folder, splits_csv=splits_csv, split=split, space=space ) if len(samples) < world: - raise RuntimeError(f"{len(samples)} atlas studies cannot supply {world} ranks") + raise RuntimeError(f"{len(samples)} {space} studies cannot supply {world} ranks") fingerprint = raw_fingerprint( samples, world, @@ -233,12 +269,13 @@ def distributed_raw_assignment( target_spacing, target_shape, posterior_shift_mm, + space, ) assignments = balanced_sample_assignment(samples, world) if max_studies_per_rank: assignments = [items[:max_studies_per_rank] for items in assignments] if any(not items for items in assignments): - raise RuntimeError("Raw atlas assignment left at least one distributed rank empty") + raise RuntimeError("Raw assignment left at least one distributed rank empty") assignment_dir.mkdir(parents=True, exist_ok=True) for assigned_rank, items in enumerate(assignments): destination = assignment_dir / f"rank_{assigned_rank:04d}.json" @@ -246,7 +283,8 @@ def distributed_raw_assignment( temporary.write_text(json.dumps(items, separators=(",", ":"))) os.replace(temporary, destination) metadata = { - "format": "mrrate_atlas_raw_assignment_v1", + "format": "mrrate_aligned_raw_assignment_v1", + "space": space, "world_size": world, "dataset_fingerprint": fingerprint, "split": split, @@ -257,13 +295,13 @@ def distributed_raw_assignment( if dist.is_initialized(): dist.barrier() metadata = json.loads(metadata_path.read_text()) - if metadata.get("format") != "mrrate_atlas_raw_assignment_v1": - raise RuntimeError(f"Invalid raw atlas assignment metadata: {metadata_path}") + if metadata.get("format") != "mrrate_aligned_raw_assignment_v1" or metadata.get("space") != space: + raise RuntimeError(f"Invalid or other-space raw assignment metadata: {metadata_path}") if int(metadata.get("world_size", -1)) != world: - raise RuntimeError("Raw atlas assignment world size changed during startup") + raise RuntimeError("Raw assignment world size changed during startup") assigned = json.loads((assignment_dir / f"rank_{rank:04d}.json").read_text()) if not assigned: - raise RuntimeError(f"Rank {rank} received no raw atlas studies") + raise RuntimeError(f"Rank {rank} received no raw {space} studies") return assigned, str(metadata["dataset_fingerprint"]) @@ -358,7 +396,7 @@ def save_checkpoint( rng_by_rank = [local_rng] if rank == 0: state = { - "format": "mrrate_atlas_dinov3d_full_v1", + "format": CHECKPOINT_FORMAT, "step": step, "stage": args.stage, "world_size": dist.get_world_size() if dist.is_initialized() else 1, @@ -395,16 +433,14 @@ def load_checkpoint( current_stage: str, dataset_fingerprint: str | None, world: int, + current_space: str = DEFAULT_SPACE, ) -> int: if not path.exists(): if rank == 0: print(f"[resume] {path} not found; starting from scratch", flush=True) return 0 state = torch.load(path, map_location="cpu", weights_only=False) - if state.get("format") != "mrrate_atlas_dinov3d_full_v1": - raise RuntimeError( - f"Checkpoint format {state.get('format')!r} is not atlas-space MR DINO" - ) + check_checkpoint_space(state.get("format"), state.get("args", {}), current_space, CHECKPOINT_FORMAT) if int(state.get("world_size", world)) != world: raise RuntimeError("Exact DDP resume requires the checkpoint's original world size") saved_fingerprint = state.get("args", {}).get("dataset_fingerprint") @@ -505,6 +541,7 @@ def main() -> int: target_spacing=tuple(args.target_spacing), target_shape=tuple(args.target_shape), posterior_shift_mm=args.posterior_shift_mm, + space=args.space, ) dataset_kwargs = { "data_folder": args.data_folder, @@ -519,7 +556,7 @@ def main() -> int: local_shape=tuple(args.local_shape) if args.local_shape else base_crop_spec.local_shape, local_crops=args.local_crops, ) - dataset = MRAtlasDINO3DDataset( + dataset = MRAlignedDINO3DDataset( **dataset_kwargs, splits_csv=args.splits_csv, split=args.split, @@ -531,6 +568,7 @@ def main() -> int: cross_sequence_probability=args.cross_sequence_probability, candidate_trials=args.candidate_trials, seed=args.seed, + global_overlap=args.global_overlap, ) local_instances = torch.tensor(len(dataset), device=device, dtype=torch.long) total_instances = local_instances.clone() @@ -548,6 +586,7 @@ def main() -> int: collate_dino3d, patch_size=(2, 16, 16), mask_ratio=(args.mask_min, args.mask_max), + cross_view_sequences=args.cross_view_sequences if args.cross_view_weight > 0 else "none", ) loader = DataLoader( dataset, @@ -566,6 +605,11 @@ def main() -> int: backbone, prototypes=args.prototypes, head_hidden_dim=args.head_hidden_dim, + dino_bottleneck_dim=args.bottleneck_dim, + ibot_bottleneck_dim=args.bottleneck_dim, + normalize_prototypes=args.normalize_prototypes, + position_bins=None if tuple(args.position_bins) == (1, 1, 1) else tuple(args.position_bins), + cross_view_weight=args.cross_view_weight, loss_weights=LossWeights( dino=1.0, ibot=args.ibot_weight, @@ -582,7 +626,11 @@ def main() -> int: learner.student.backbone.blocks[i] = checkpoint_wrapper(block) learner.to(device) trainable = [p for p in learner.parameters() if p.requires_grad] - optimizer = torch.optim.AdamW(trainable, lr=args.lr, betas=(0.9, 0.999), weight_decay=args.weight_decay) + optimizer = make_optimizer( + learner.student, lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.999), + layerwise_decay=args.layerwise_decay, patch_embed_lr_mult=args.patch_embed_lr_mult, + n_blocks=len(learner.student.backbone.blocks), fused=device.type == "cuda", + ) if world > 1: learner = DDP( learner, @@ -595,7 +643,7 @@ def main() -> int: start = load_checkpoint( checkpoint_path(output, args.resume), learner, optimizer, sampler, device, rank, - args.stage, args.dataset_fingerprint, world, + args.stage, args.dataset_fingerprint, world, args.space, ) if checkpoint_path(output, args.resume) else 0 consumed_samples = sampler.offset # Cross-stage resume fixes the new Gram target to the incoming EMA teacher. @@ -657,12 +705,15 @@ def main() -> int: if not finite.item(): raise FloatingPointError(f"Non-finite distributed loss at step {step}: {loss.detach()}") backward_loss.backward() - grad_norm = torch.nn.utils.clip_grad_norm_(trainable, 3.0) + if step < args.freeze_last_layer_steps: + cancel_last_layer_gradients(unwrap(learner).student) # DINO/DINOv2 freeze_last_layer + grad_norm = torch.nn.utils.clip_grad_norm_(trainable, args.clip_grad) optimizer.step() unwrap(learner).update_teacher(momentum) if (step + 1) % args.log_every == 0: - values = torch.stack([metrics[k].float() for k in ("loss", "dino", "ibot", "koleo", "gram")]) + metric_keys = sorted(metrics) + values = torch.stack([metrics[k].float() for k in metric_keys]) if dist.is_initialized(): dist.all_reduce(values) values /= world @@ -676,11 +727,7 @@ def main() -> int: "sequence_epoch": ( (step + 1) * global_batch / int(total_instances) ), - "loss": float(values[0]), - "dino": float(values[1]), - "ibot": float(values[2]), - "koleo": float(values[3]), - "gram": float(values[4]), + **{k: float(v) for k, v in zip(metric_keys, values)}, "lr": lr, "weight_decay": wd, "teacher_momentum": momentum, diff --git a/dino-pretraining/mr_dino/train_fsdp.py b/dino-pretraining/mr_dino/train_fsdp.py index f105e42..aab63f0 100644 --- a/dino-pretraining/mr_dino/train_fsdp.py +++ b/dino-pretraining/mr_dino/train_fsdp.py @@ -1,4 +1,4 @@ -"""FSDP2 training for the literal DINOv3 ViT-7B on atlas-registered MR-RATE.""" +"""FSDP2 3-D DINOv3 training (ViT-L default) on aligned (coreg/atlas) MR-RATE.""" from __future__ import annotations @@ -26,16 +26,20 @@ from dinov3.layers.fp8_linear import convert_linears_to_fp8 from .data import ( - MRAtlasDINO3DDataset, + ALIGNED_SPACES, + DEFAULT_SPACE, + MRAlignedDINO3DDataset, InfiniteStudySampler, collate_dino3d, stage_crop_spec, ) from .fp8 import enable_fsdp_mixed_precision_fp8 from .objective import DINO3DLearner, LossWeights +from .recipe import cancel_last_layer_gradients, make_optimizer, missing_frozen_optimizer_keys from .train_ddp import ( StopController, balanced_file_assignment, + check_checkpoint_space, build_backbone, cache_fingerprint, cosine, @@ -46,29 +50,34 @@ ) +CHECKPOINT_FORMAT = "mrrate_aligned_dinov3d_fsdp2_v1" + + def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description=__doc__) source = p.add_mutually_exclusive_group(required=True) source.add_argument( "--data-folder", - help="MR-RATE-atlas NIfTI tree; uses the same live loader as previous MIL training", + help="MR-RATE-coreg (or -atlas) NIfTI tree; uses the same live loader as previous MIL training", ) source.add_argument( "--preprocessed-dir", - help="Optional MR-RATE volume cache root containing atlas_space/*.npz", + help="Optional MR-RATE volume cache root containing /*.npz", ) p.add_argument("--splits-csv") p.add_argument("--split", default="train") - p.add_argument("--space", default="atlas_space", choices=("atlas_space",)) + p.add_argument("--space", default=DEFAULT_SPACE, choices=ALIGNED_SPACES, + help="coreg_space (default): sequences on the center T1's native grid; atlas_space: MNI 1 mm box") p.add_argument("--output-dir", required=True) p.add_argument("--stage", choices=("pretrain", "gram", "highres"), default="pretrain") - p.add_argument("--arch", choices=("tiny", "7b"), default="7b") - p.add_argument("--steps", type=int, default=100_000) - p.add_argument("--batch-size", type=int, default=1) + p.add_argument("--arch", choices=("tiny", "large", "hplus", "7b"), default="hplus", + help="hplus = ViT-H+ 0.84B (MR default); large = ViT-L 0.30B (FORA CT); 7b kept for reference") + p.add_argument("--steps", type=int, default=50_000) + p.add_argument("--batch-size", type=int, default=8) p.add_argument( "--grad-accum-steps", type=int, - default=4, + default=1, help="Fixed study microbatches per optimizer update on every rank", ) p.add_argument("--workers", type=int, default=2) @@ -81,18 +90,18 @@ def parse_args() -> argparse.Namespace: p.add_argument("--target-shape", type=int, nargs=3, default=(256, 384, 384)) p.add_argument("--posterior-shift-mm", type=float, default=15.0) p.add_argument("--seed", type=int, default=3407) - p.add_argument("--dino-prototypes", type=int, default=262_144) - p.add_argument("--ibot-prototypes", type=int, default=98_304) - p.add_argument("--dino-head-hidden-dim", type=int, default=8192) - p.add_argument("--ibot-head-hidden-dim", type=int, default=4096) - p.add_argument("--dino-bottleneck-dim", type=int, default=512) - p.add_argument("--ibot-bottleneck-dim", type=int, default=384) - p.add_argument("--lr", type=float, default=5e-5, help="Peak at global batch 1024 before sqrt scaling") + p.add_argument("--dino-prototypes", type=int, default=16_384) + p.add_argument("--ibot-prototypes", type=int, default=16_384) + p.add_argument("--dino-head-hidden-dim", type=int, default=2048) + p.add_argument("--ibot-head-hidden-dim", type=int, default=2048) + p.add_argument("--dino-bottleneck-dim", type=int, default=256) + p.add_argument("--ibot-bottleneck-dim", type=int, default=256) + p.add_argument("--lr", type=float, default=1e-3, help="Peak at global batch 1024 before sqrt scaling (DINOv2 ViT-L)") p.add_argument("--min-lr", type=float, default=1e-6) - p.add_argument("--warmup-steps", type=int, default=10_000) + p.add_argument("--warmup-steps", type=int, default=2_500) p.add_argument("--weight-decay", type=float, default=0.04) - p.add_argument("--weight-decay-end", type=float, default=0.04) - p.add_argument("--teacher-momentum", type=float, default=0.994) + p.add_argument("--weight-decay-end", type=float, default=0.4) + p.add_argument("--teacher-momentum", type=float, default=0.992) p.add_argument("--teacher-temperature", type=float, default=0.07) p.add_argument("--teacher-warmup-temperature", type=float, default=0.04) p.add_argument("--teacher-warmup-steps", type=int, default=10_000) @@ -103,18 +112,28 @@ def parse_args() -> argparse.Namespace: p.add_argument("--gram-weight", type=float, default=1.0) p.add_argument("--gram-max-tokens", type=int, default=1024) p.add_argument("--ibot-loss-chunk-size", type=int, default=1024) - p.add_argument("--cross-sequence-probability", type=float, default=0.75) + p.add_argument("--cross-sequence-probability", type=float, default=0.25) p.add_argument("--candidate-trials", type=int, default=12) - p.add_argument("--clip-grad", type=float, default=30.0) - p.add_argument("--save-every", type=int, default=250) - p.add_argument("--keep-checkpoints", type=int, default=3) + p.add_argument("--clip-grad", type=float, default=3.0) + p.add_argument("--save-every", type=int, default=500) + p.add_argument("--keep-checkpoints", type=int, default=6) p.add_argument("--log-every", type=int, default=10) p.add_argument("--resume", default="latest", help="latest, none, or a DCP checkpoint directory") - p.add_argument("--activation-checkpointing", action=argparse.BooleanOptionalAction, default=True) + p.add_argument("--activation-checkpointing", action=argparse.BooleanOptionalAction, default=False) p.add_argument("--selective-checkpointing", action=argparse.BooleanOptionalAction, default=True) p.add_argument("--compile", action=argparse.BooleanOptionalAction, default=True) - p.add_argument("--fp8", action=argparse.BooleanOptionalAction, default=True) + p.add_argument("--fp8", action=argparse.BooleanOptionalAction, default=False) p.add_argument("--max-studies", type=int, default=None, help="Smoke only: cap studies per rank") + # FORA recipe fixes (see mr_dino/recipe.py) + p.add_argument("--global-overlap", type=float, default=0.25, help="Minimum per-axis overlap of the two global crops") + p.add_argument("--normalize-prototypes", action=argparse.BooleanOptionalAction, default=True) + p.add_argument("--position-bins", type=int, nargs=3, default=(2, 2, 2), help="Spatial Sinkhorn bins for patch targets; 1 1 1 disables") + p.add_argument("--cross-view-weight", type=float, default=1.0, help="Cross-crop twin-token patch objective (0 disables)") + p.add_argument("--cross-view-sequences", choices=("same", "any", "none"), default="same") + p.add_argument("--layerwise-decay", type=float, default=0.9) + p.add_argument("--patch-embed-lr-mult", type=float, default=0.2) + p.add_argument("--freeze-last-layer-steps", type=int, default=1250) + p.add_argument("--adamw-beta2", type=float, default=0.999) args = p.parse_args() # A running Slurm allocation can cleanly restart the worker step without # relinquishing its nodes. Keep schedule changes in a small, auditable @@ -245,6 +264,9 @@ def build_learner(args: argparse.Namespace, device: torch.device) -> DINO3DLearn ibot_loss_chunk_size=args.ibot_loss_chunk_size, with_gram_anchor=args.stage in {"gram", "highres"}, initialize_weights=False, + normalize_prototypes=args.normalize_prototypes, + position_bins=None if tuple(args.position_bins) == (1, 1, 1) else tuple(args.position_bins), + cross_view_weight=args.cross_view_weight, ) if args.activation_checkpointing: wrapper = _checkpoint_wrapper(args.selective_checkpointing) @@ -272,15 +294,19 @@ def checkpoint_root(output: Path) -> Path: def resolve_checkpoint(output: Path, resume: str) -> Path | None: + """``latest``, ``none``, or an explicit checkpoint (normally the previous stage's last one). + + Once this stage has written its own complete checkpoint, that one wins over an explicit + path: a requeued gram/highres job keeps RESUME= in its environment and + must continue, not restart from the previous stage. + """ if not resume or resume.lower() == "none": return None - if resume != "latest": - return Path(resume) marker = checkpoint_root(output) / "latest.txt" - if not marker.exists(): - return None - path = checkpoint_root(output) / marker.read_text().strip() - return path if (path / "COMPLETE").exists() else None + own = checkpoint_root(output) / marker.read_text().strip() if marker.exists() else None + if own is not None and (own / "COMPLETE").exists(): + return own + return None if resume == "latest" else Path(resume) def runtime_state(sampler_offset: int) -> dict: @@ -310,7 +336,7 @@ def save_checkpoint( os.replace(tmp, path / f"runtime_rank{rank:04d}.pt") if rank == 0: (path / "metadata.json").write_text(json.dumps({ - "format": "mrrate_atlas_dinov3d_fsdp2_v1", + "format": CHECKPOINT_FORMAT, "step": step, "stage": args.stage, "sampler_offset_per_rank": sampler_offset, @@ -346,10 +372,7 @@ def load_checkpoint( if not (path / "COMPLETE").exists(): raise RuntimeError(f"Incomplete checkpoint: {path}") metadata = json.loads((path / "metadata.json").read_text()) - if metadata.get("format") != "mrrate_atlas_dinov3d_fsdp2_v1": - raise RuntimeError( - f"Checkpoint format {metadata.get('format')!r} is not atlas-space MR DINO" - ) + check_checkpoint_space(metadata.get("format"), metadata.get("args", {}), args.space, CHECKPOINT_FORMAT) saved_world = int(metadata.get("world_size", dist.get_world_size())) if saved_world != dist.get_world_size(): raise RuntimeError( @@ -359,15 +382,21 @@ def load_checkpoint( saved_fingerprint = metadata.get("args", {}).get("dataset_fingerprint") if saved_fingerprint and saved_fingerprint != args.dataset_fingerprint: raise RuntimeError( - "Atlas-space training cache/split changed since the checkpoint; " + "Training cache/split/space changed since the checkpoint; " "refusing an inexact sampler resume" ) source_stage = metadata["stage"] + same_stage = source_stage == args.stage + # checkpoints written during freeze_last_layer have no optimizer state for the frozen prototype layers + frozen_missing = missing_frozen_optimizer_keys(path, learner, optimizer) if same_stage else [] + if frozen_missing and rank == 0: + print(f"[resume] {len(frozen_missing)} optimizer-state keys of frozen last_layer params absent " + "(saved during freeze_last_layer); they start fresh", flush=True) loaded_step = int(dcp_load( path, model=learner, optimizer=optimizer, - strict_loading=(source_stage == args.stage), + strict_loading=same_stage and not frozen_missing, )) state = torch.load(path / f"runtime_rank{rank:04d}.pt", map_location="cpu", weights_only=False) sampler.set_offset(int(state["sampler_offset"])) @@ -438,6 +467,7 @@ def main() -> int: target_spacing=tuple(args.target_spacing), target_shape=tuple(args.target_shape), posterior_shift_mm=args.posterior_shift_mm, + space=args.space, ) dataset_kwargs = { "data_folder": args.data_folder, @@ -451,7 +481,7 @@ def main() -> int: local_shape=tuple(args.local_shape) if args.local_shape else base_crop_spec.local_shape, local_crops=args.local_crops, ) - dataset = MRAtlasDINO3DDataset( + dataset = MRAlignedDINO3DDataset( **dataset_kwargs, splits_csv=args.splits_csv, split=args.split, @@ -463,6 +493,7 @@ def main() -> int: cross_sequence_probability=args.cross_sequence_probability, candidate_trials=args.candidate_trials, seed=args.seed, + global_overlap=args.global_overlap, ) local_instances = torch.tensor(len(dataset), device=device, dtype=torch.long) total_instances = local_instances.clone() @@ -477,13 +508,14 @@ def main() -> int: ) learner = build_learner(args, device) - trainable = list(learner.student.parameters()) - optimizer = torch.optim.AdamW( - trainable, + optimizer = make_optimizer( + learner.student, lr=args.lr, - betas=(0.9, 0.99), weight_decay=args.weight_decay, - fused=True, + betas=(0.9, args.adamw_beta2), + layerwise_decay=args.layerwise_decay, + patch_embed_lr_mult=args.patch_embed_lr_mult, + n_blocks=len(learner.student.backbone.blocks), ) start = load_checkpoint(resolve_checkpoint(output, args.resume), learner, optimizer, sampler, args, rank) transition_step = int(getattr(args, "warmup_transition_step", 0)) @@ -497,6 +529,7 @@ def main() -> int: collate_dino3d, patch_size=(2, 16, 16), mask_ratio=(args.mask_min, args.mask_max), + cross_view_sequences=args.cross_view_sequences if args.cross_view_weight > 0 else "none", ) loader_generator = torch.Generator().manual_seed(args.seed + rank * 104729) loader = DataLoader( @@ -566,7 +599,7 @@ def main() -> int: group["weight_decay"] = wd optimizer.zero_grad(set_to_none=True) - metric_accumulator = torch.zeros(5, device=device, dtype=torch.float32) + metric_accumulator = None for microstep in range(args.grad_accum_steps): batch = next(iterator) consumed_samples += args.batch_size @@ -596,9 +629,11 @@ def main() -> int: f"{backward_loss.detach()}; metrics={metric_debug}" ) backward_loss.backward() - metric_accumulator += torch.stack( - [metrics[k].float() for k in ("loss", "dino", "ibot", "koleo", "gram")] - ) + metric_keys = sorted(metrics) + values_now = torch.stack([metrics[k].float() for k in metric_keys]) + metric_accumulator = values_now if metric_accumulator is None else metric_accumulator + values_now + if step < args.freeze_last_layer_steps: + cancel_last_layer_gradients(learner.student) # DINO/DINOv2 freeze_last_layer grad_norms = [ torch.nn.utils.clip_grad_norm_(module.parameters(), args.clip_grad) for module in (learner.student.backbone, learner.student.dino_head, learner.student.ibot_head) @@ -621,11 +656,7 @@ def main() -> int: / int(total_instances) ), "grad_accum_steps": args.grad_accum_steps, - "loss": float(values[0]), - "dino": float(values[1]), - "ibot": float(values[2]), - "koleo": float(values[3]), - "gram": float(values[4]), + **{k: float(v) for k, v in zip(metric_keys, values)}, "lr": lr, "weight_decay": wd, "teacher_momentum": momentum, diff --git a/dino-pretraining/scripts/probe_checkpoint.sbatch b/dino-pretraining/scripts/probe_checkpoint.sbatch new file mode 100755 index 0000000..7a76b0a --- /dev/null +++ b/dino-pretraining/scripts/probe_checkpoint.sbatch @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +#SBATCH --job-name=mr_dino_probe +#SBATCH --partition=h200 +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:h200:4 +#SBATCH --cpus-per-task=128 +#SBATCH --time=06:00:00 +#SBATCH --output=mrdino_%x_%j.out +#SBATCH --error=mrdino_%x_%j.err +# Usage: DATA_FOLDER= sbatch probe_checkpoint.sbatch CHECKPOINT OUT_DIR [extra extract args] +set -euo pipefail +SCRIPT_DIR=${SCRIPT_DIR:-"$SLURM_SUBMIT_DIR/scripts"} +[[ -x "$SCRIPT_DIR/probe_node.sh" ]] || SCRIPT_DIR="$SLURM_SUBMIT_DIR/dino-pretraining/scripts" +exec "$SCRIPT_DIR/probe_node.sh" "$@" diff --git a/dino-pretraining/scripts/probe_node.sh b/dino-pretraining/scripts/probe_node.sh new file mode 100755 index 0000000..fe10e77 --- /dev/null +++ b/dino-pretraining/scripts/probe_node.sh @@ -0,0 +1,56 @@ +#!/usr/bin/env bash +# Score one MR DINO checkpoint on THIS node with the validation MIL probe: +# extract (one worker per GPU) -> patient-level K folds in parallel -> out-of-fold summary. +# Usage: DATA_FOLDER= probe_node.sh CHECKPOINT OUT_DIR [extra extract args] +# CHECKPOINT: DCP step directory or the teacher .pt written by `python -m mr_dino.probe export`. +# Tokens go to node-local /tmp (FEATURE_DIR) and are deleted afterwards unless KEEP_FEATURES=1. +# Inside a running training allocation: srun --overlap --jobid JOB -N1 -n1 -w NODE probe_node.sh ... +set -euo pipefail + +CKPT=${1:?checkpoint}; OUT=${2:?output dir}; shift 2 +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CODE=${CODE:-"$(dirname "$SCRIPT_DIR")"} +DINOV3_ROOT=${DINOV3_ROOT:-/hnvme/workspace/b180dc51-sezgin/dinov3} +SIF=${SIF:-/hnvme/workspace/b180dc51-sezgin/mrrate.sif} +EXTRA_PIP=${EXTRA_PIP:-/hnvme/workspace/b180dc51-sezgin/extra-pip} +: "${DATA_FOLDER:?Set DATA_FOLDER to the extracted MR-RATE tree (batchXX//coreg_img)}" +GPUS=${GPUS_PER_NODE:-4} +FOLDS=${FOLDS:-5} +FEAT=${FEATURE_DIR:-/tmp/mrdino_probe_${SLURM_JOB_ID:-$$}_$(basename "$CKPT" .pt)} +LABEL_ARGS=() +[[ -n "${LABELS_CSV:-}" ]] && LABEL_ARGS+=(--labels-csv "$LABELS_CSV") +[[ -n "${SPLITS_CSV:-}" ]] && LABEL_ARGS+=(--splits-csv "$SPLITS_CSV") +export PYTHONPATH="${EXTRA_PIP}:${DINOV3_ROOT}:${CODE}:${PYTHONPATH:-}" +export OMP_NUM_THREADS=${OMP_NUM_THREADS:-8} + +probe() { + singularity exec --nv -B /hnvme/workspace:/hnvme/workspace,/tmp:/tmp "$SIF" \ + python3 -m mr_dino.probe "$@" +} +wait_all() { # fail if any background worker failed + local status=0 + for pid in "$@"; do wait "$pid" || status=1; done + return $status +} + +mkdir -p "$OUT" "$FEAT" +echo "[probe_node] $(date '+%F %T') checkpoint=$CKPT features=$FEAT out=$OUT host=$(hostname)" +pids=() +for g in $(seq 0 $((GPUS - 1))); do + CUDA_VISIBLE_DEVICES=$g RANK=$g WORLD_SIZE=$GPUS probe extract --checkpoint "$CKPT" \ + --data-folder "$DATA_FOLDER" --features-dir "$FEAT" "${LABEL_ARGS[@]}" "$@" \ + > "$OUT/extract_rank$g.log" 2>&1 & + pids+=($!) +done +wait_all "${pids[@]}" || { echo "[probe_node] extraction failed; see $OUT/extract_rank*.log" >&2; exit 1; } +echo "[probe_node] $(date '+%F %T') extraction done ($(du -sh "$FEAT" | cut -f1))" + +pids=() +for f in $(seq 0 $((FOLDS - 1))); do + CUDA_VISIBLE_DEVICES=$((f % GPUS)) probe cv --features-dir "$FEAT" --out-dir "$OUT" \ + --fold "$f" --folds "$FOLDS" "${LABEL_ARGS[@]}" > "$OUT/cv_fold$f.log" 2>&1 & + pids+=($!) +done +wait_all "${pids[@]}" || { echo "[probe_node] CV failed; see $OUT/cv_fold*.log" >&2; exit 1; } +probe summarize --features-dir "$FEAT" --out-dir "$OUT" --folds "$FOLDS" "${LABEL_ARGS[@]}" | tee "$OUT/summary.txt" +[[ "${KEEP_FEATURES:-0}" == 1 ]] || rm -rf "$FEAT" diff --git a/dino-pretraining/scripts/probe_watch.sh b/dino-pretraining/scripts/probe_watch.sh new file mode 100755 index 0000000..ba8dbcf --- /dev/null +++ b/dino-pretraining/scripts/probe_watch.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +# Probe chosen steps of a running MR DINO stage as soon as their checkpoints appear. +# Usage (login node, e.g. under nohup): DATA_FOLDER= probe_watch.sh RUN_DIR STEP [STEP ...] +# RUN_DIR is a stage output (…/mrdino3d_hplus/pretrain). For each step the EMA-teacher backbone is +# exported at once to RUN_DIR/probe/step_N/teacher.pt (checkpoint rotation cannot delete it), then +# probe_checkpoint.sbatch is submitted. Results: RUN_DIR/probe/step_N/results.json + summary.txt. +set -euo pipefail +RUN=${1:?run dir}; shift +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CODE=${CODE:-"$(dirname "$SCRIPT_DIR")"} +SIF=${SIF:-/hnvme/workspace/b180dc51-sezgin/mrrate.sif} +EXTRA_PIP=${EXTRA_PIP:-/hnvme/workspace/b180dc51-sezgin/extra-pip} +DINOV3_ROOT=${DINOV3_ROOT:-/hnvme/workspace/b180dc51-sezgin/dinov3} +: "${DATA_FOLDER:?Set DATA_FOLDER}" +export DATA_FOLDER PYTHONPATH="${EXTRA_PIP}:${DINOV3_ROOT}:${CODE}:${PYTHONPATH:-}" +pending=("$@") +while ((${#pending[@]})); do + left=() + for step in "${pending[@]}"; do + ckpt=$(printf '%s/checkpoints/step_%08d' "$RUN" "$step") + out="$RUN/probe/step_$step" + if [[ -f "$ckpt/COMPLETE" ]]; then + mkdir -p "$out" + singularity exec -B /hnvme/workspace:/hnvme/workspace "$SIF" \ + python3 -m mr_dino.probe export --checkpoint "$ckpt" --out "$out/teacher.pt" + job=$(cd "$CODE" && sbatch --parsable --chdir="$CODE" --output="$out/slurm_%j.out" \ + --error="$out/slurm_%j.err" "$SCRIPT_DIR/probe_checkpoint.sbatch" "$out/teacher.pt" "$out") + echo "[probe_watch] $(date '+%F %T') step $step exported, probe job $job" + else + left+=("$step") + fi + done + pending=("${left[@]+"${left[@]}"}") + ((${#pending[@]})) && sleep 300 +done +echo "[probe_watch] all steps submitted" diff --git a/dino-pretraining/scripts/smoke_dummy.sbatch b/dino-pretraining/scripts/smoke_dummy.sbatch index 59081c5..3119ece 100755 --- a/dino-pretraining/scripts/smoke_dummy.sbatch +++ b/dino-pretraining/scripts/smoke_dummy.sbatch @@ -28,14 +28,14 @@ EXTRA_PIP=${EXTRA_PIP:-/hnvme/workspace/b180dc51-sezgin/extra-pip} export PYTHONPATH="${EXTRA_PIP}:${PYTHONPATH:-}" singularity exec -B /hnvme/workspace:/hnvme/workspace,/tmp:/tmp \ "$SIF" \ - python3 "$CODE/tests/make_dummy_atlas_dataset.py" \ + python3 "$CODE/tests/make_dummy_aligned_dataset.py" \ --out "$DUMMY" --studies 8 --shape 8 32 32 head_node=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1) MASTER_ADDR=$(getent ahostsv4 "$head_node" | awk '{print $1; exit}') export MASTER_ADDR MASTER_PORT=${MASTER_PORT:-29672} COMMON_ARGS=( - --data-folder "$DUMMY" --space atlas_space --output-dir "$OUTPUT" \ + --data-folder "$DUMMY" --space coreg_space --output-dir "$OUTPUT" \ --stage pretrain --arch tiny --batch-size 1 --grad-accum-steps 1 \ --workers 0 --max-studies 2 --local-crops 2 \ --target-shape 8 32 32 --target-spacing 1.0 0.5 0.5 --posterior-shift-mm 0 \ @@ -57,4 +57,32 @@ srun --nodes=1 --ntasks=1 --kill-on-bad-exit=1 \ test -f "$OUTPUT/checkpoints/step_00000002/COMPLETE" grep -q '"step": 2' "$OUTPUT/metrics.jsonl" -echo "Synthetic raw-atlas MR DINO FSDP2 integration and resume passed: $OUTPUT" +echo "Synthetic raw-coreg MR DINO FSDP2 integration and resume passed: $OUTPUT" + +# Stage transitions with the phase-2/3 launcher settings (fixed teacher temperature, no freeze). +STAGE_FIXED=(--teacher-temperature 0.07 --teacher-warmup-temperature 0.07 --teacher-warmup-steps 1 + --teacher-momentum 0.996 --freeze-last-layer-steps 0 --warmup-steps 1) +srun --nodes=1 --ntasks=1 --kill-on-bad-exit=1 "$SCRIPT_DIR/train_node.sh" "${COMMON_ARGS[@]}" \ + --output-dir "$OUTPUT/gram" --stage gram --gram-weight 1.0 --steps 2 \ + --resume "$OUTPUT/checkpoints/step_00000002" "${STAGE_FIXED[@]}" +test -f "$OUTPUT/gram/checkpoints/step_00000002/COMPLETE" +srun --nodes=1 --ntasks=1 --kill-on-bad-exit=1 "$SCRIPT_DIR/train_node.sh" "${COMMON_ARGS[@]}" \ + --output-dir "$OUTPUT/highres" --stage highres --gram-weight 1.5 --steps 2 \ + --resume "$OUTPUT/gram/checkpoints/step_00000002" "${STAGE_FIXED[@]}" +test -f "$OUTPUT/highres/checkpoints/step_00000002/COMPLETE" +echo "Stage transitions pretrain -> gram -> highres passed" + +# Validation probe on the step-2 checkpoint: all dummy studies labelled "val", 2 folds. +python3 - "$DUMMY" <<'PY' +import csv, pathlib, sys +root = pathlib.Path(sys.argv[1]) +uids = sorted(p.name for p in (root / "batch00").iterdir()) +with open(root / "labels.csv", "w", newline="") as f: + w = csv.writer(f); w.writerow(["study_uid", "a", "b"]); w.writerows([u, i % 2, (i // 2) % 2] for i, u in enumerate(uids)) +with open(root / "splits.csv", "w", newline="") as f: + w = csv.writer(f); w.writerow(["study_uid", "patient_uid", "split"]); w.writerows([u, u, "val"] for u in uids) +PY +DATA_FOLDER="$DUMMY" LABELS_CSV="$DUMMY/labels.csv" SPLITS_CSV="$DUMMY/splits.csv" FOLDS=2 \ + "$SCRIPT_DIR/probe_node.sh" "$OUTPUT/checkpoints/step_00000002" "$OUTPUT/probe" --tile 8 32 32 --workers 0 +test -f "$OUTPUT/probe/results.json" +echo "Validation probe smoke passed: $OUTPUT/probe/summary.txt" diff --git a/dino-pretraining/scripts/train_32n_7b.sbatch b/dino-pretraining/scripts/train_32n_7b.sbatch index d7dee8b..8a41253 100755 --- a/dino-pretraining/scripts/train_32n_7b.sbatch +++ b/dino-pretraining/scripts/train_32n_7b.sbatch @@ -1,4 +1,5 @@ #!/usr/bin/env bash +# LEGACY recipe (DINOv3-7B continuation settings). The FORA CT run with these settings failed (flat prototype targets, 1k warmup, last-block position leakage); use train_mrdino.sbatch. #SBATCH --job-name=mr_dino3_7b #SBATCH --partition=h200 #SBATCH --nodes=32 diff --git a/dino-pretraining/scripts/train_32n_hplus.sbatch b/dino-pretraining/scripts/train_32n_hplus.sbatch index 2daf2b5..3b2c1af 100755 --- a/dino-pretraining/scripts/train_32n_hplus.sbatch +++ b/dino-pretraining/scripts/train_32n_hplus.sbatch @@ -1,4 +1,5 @@ #!/usr/bin/env bash +# LEGACY recipe (DINOv3-7B continuation settings). The FORA CT run with these settings failed (flat prototype targets, 1k warmup, last-block position leakage); use train_mrdino.sbatch. #SBATCH --job-name=mr_dino3_hplus #SBATCH --partition=h200 #SBATCH --nodes=32 diff --git a/dino-pretraining/scripts/train_mrdino.sbatch b/dino-pretraining/scripts/train_mrdino.sbatch new file mode 100644 index 0000000..edd3f2f --- /dev/null +++ b/dino-pretraining/scripts/train_mrdino.sbatch @@ -0,0 +1,102 @@ +#!/usr/bin/env bash +#SBATCH --job-name=mr_dino3_hplus +#SBATCH --partition=h200 +#SBATCH --nodes=4 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:h200:4 +#SBATCH --cpus-per-task=128 +#SBATCH --exclusive +#SBATCH --exclude=h21-24,h23-24,h24-17,h23-18 +#SBATCH --time=1-00:00:00 +#SBATCH --requeue +#SBATCH --signal=B:USR1@1200 +#SBATCH --output=mrdino_%x_%j.out +#SBATCH --error=mrdino_%x_%j.err +#SBATCH --mail-type=ALL +#SBATCH --mail-user=sezgin.er@dqbm.uzh.ch +set -uo pipefail + +if [[ -z "${CODE:-}" ]]; then + if [[ -d "$SLURM_SUBMIT_DIR/mr_dino" ]]; then + CODE="$SLURM_SUBMIT_DIR" + elif [[ -d "$SLURM_SUBMIT_DIR/dino-pretraining/mr_dino" ]]; then + CODE="$SLURM_SUBMIT_DIR/dino-pretraining" + else + echo "Cannot locate dino-pretraining; set CODE explicitly" >&2 + exit 2 + fi +fi +SCRIPT_DIR="$CODE/scripts" +: "${DATA_FOLDER:?Set DATA_FOLDER to the extracted MR-RATE-coreg NIfTI root (batchXX//coreg_img)}" +STAGE=${STAGE:-pretrain} +OUTPUT=${OUTPUT:-"$PWD/mrdino3d_${ARCH:-hplus}/$STAGE"} +mkdir -p "$OUTPUT" +rm -f "$OUTPUT/REQUEUE_REQUESTED" + +head_node=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1) +MASTER_ADDR=$(getent ahostsv4 "$head_node" | awk '{print $1; exit}') +export MASTER_ADDR MASTER_PORT=${MASTER_PORT:-29675} +end_ts=$(scontrol show job "$SLURM_JOB_ID" | sed -n 's/.*EndTime=\([0-9T:-]*\).*/\1/p' | head -1) +if [[ -n "$end_ts" ]]; then + export MRDINO_DEADLINE_EPOCH=$(( $(date -d "$end_ts" +%s) - 1200 )) +fi + +# Phases 2/3 restart the schedule at step 0, so they need their own settings (FORA CT-DINO +# phase2_gram_recipe.args / phase3_hires_recipe.args): short warmup, teacher temperature fixed at +# 0.07, EMA 0.996, no prototype freeze. Phase 3 has 4x the tokens per global crop -> batch 4 + checkpointing. +# --lr is the peak at global batch 1024 and is sqrt-scaled to the real batch. +# H+ memory/speed per H200 (1-node benchmark 943261, real crop sizes, compiled, bf16): +# phase-1 crops: batch 8 no checkpointing OOM; batch 8 + selective checkpointing 0.39 steps/s 32 GiB; +# batch 4 no checkpointing 0.95 steps/s 83 GiB. +# highres crops (Gram on): batch 2 no checkpointing 128 GiB (too tight); batch 2 + checkpointing +# 0.23 steps/s 43 GiB; batch 4 + checkpointing 0.10 steps/s 81 GiB. +# Defaults keep the FORA global batch per phase on 4 nodes: 128 (phases 1-2), 64 (phase 3). +ARCH=${ARCH:-hplus} +FIXED_TEACHER=(--teacher-temperature 0.07 --teacher-warmup-temperature 0.07 --teacher-warmup-steps 1 + --teacher-momentum 0.996 --freeze-last-layer-steps 0) +case "$STAGE" in + pretrain) # ~10 sequence-epochs at 4 nodes (batch 128); the cosine decay ends here + STEPS=${STEPS:-50000}; GRAM_WEIGHT=0; BATCH_SIZE=${BATCH_SIZE:-8}; LOCAL_CROPS=${LOCAL_CROPS:-8} + ACT_CKPT=${ACT_CKPT:-1} + STAGE_ARGS=() ;; + gram) # Gram anchoring to the incoming teacher + STEPS=${STEPS:-10000}; GRAM_WEIGHT=${GRAM_WEIGHT:-1.0}; BATCH_SIZE=${BATCH_SIZE:-8}; LOCAL_CROPS=${LOCAL_CROPS:-8} + ACT_CKPT=${ACT_CKPT:-1} + STAGE_ARGS=(--lr 5e-4 --warmup-steps 250 --weight-decay 0.1 --weight-decay-end 0.4 "${FIXED_TEACHER[@]}") ;; + highres) # 64x384x384 globals (whole axial plane, 18,432 tokens), 4 locals 32x192x192 + STEPS=${STEPS:-5000}; GRAM_WEIGHT=${GRAM_WEIGHT:-1.5}; BATCH_SIZE=${BATCH_SIZE:-4}; LOCAL_CROPS=${LOCAL_CROPS:-4} + ACT_CKPT=${ACT_CKPT:-1} + STAGE_ARGS=(--lr 4e-4 --warmup-steps 500 --weight-decay 0.04 --weight-decay-end 0.2 "${FIXED_TEACHER[@]}") ;; + *) echo "Unknown STAGE=$STAGE" >&2; exit 2 ;; +esac +if [[ "$STAGE" != pretrain && "${RESUME:-latest}" == latest && ! -f "$OUTPUT/checkpoints/latest.txt" ]]; then + echo "STAGE=$STAGE starts from the previous stage: set RESUME=/checkpoints/step_N" >&2 + exit 2 +fi +[[ "${COMPILE:-1}" == 1 ]] && COMPILE_ARGS=(--compile) || COMPILE_ARGS=(--no-compile) +[[ "$ACT_CKPT" == 1 ]] && CKPT_ARGS=(--activation-checkpointing) || CKPT_ARGS=(--no-activation-checkpointing) +[[ "${FP8:-1}" == 1 ]] && FP8_ARGS=(--fp8) || FP8_ARGS=(--no-fp8) +SPLIT_ARGS=() +if [[ -n "${SPLITS_CSV:-}" ]]; then + SPLIT_ARGS=(--splits-csv "$SPLITS_CSV") +fi + +echo "MR DINOv3 $ARCH (FORA recipe): stage=$STAGE nodes=$SLURM_NNODES output=$OUTPUT data=$DATA_FOLDER" +srun --nodes="$SLURM_NNODES" --ntasks="$SLURM_NNODES" --ntasks-per-node=1 \ + --kill-on-bad-exit=1 "$SCRIPT_DIR/train_node.sh" \ + --data-folder "$DATA_FOLDER" \ + --space "${SPACE:-coreg_space}" \ + "${SPLIT_ARGS[@]}" \ + --output-dir "$OUTPUT" --stage "$STAGE" --arch "$ARCH" --steps "$STEPS" \ + --batch-size "$BATCH_SIZE" --grad-accum-steps "${GRAD_ACCUM_STEPS:-1}" \ + --workers "${WORKERS:-8}" --local-crops "$LOCAL_CROPS" \ + --cross-sequence-probability "${CROSS_SEQUENCE_PROBABILITY:-0.25}" \ + --cross-view-sequences "${CROSS_VIEW_SEQUENCES:-same}" \ + --resume "${RESUME:-latest}" --gram-weight "$GRAM_WEIGHT" \ + "${COMPILE_ARGS[@]}" --no-fp8 "${CKPT_ARGS[@]}" "${STAGE_ARGS[@]}" +rc=$? +if [[ -f "$OUTPUT/REQUEUE_REQUESTED" ]]; then + scontrol requeue "$SLURM_JOB_ID" + exit 0 +fi +exit "$rc" diff --git a/dino-pretraining/tests/make_dummy_atlas_dataset.py b/dino-pretraining/tests/make_dummy_aligned_dataset.py similarity index 82% rename from dino-pretraining/tests/make_dummy_atlas_dataset.py rename to dino-pretraining/tests/make_dummy_aligned_dataset.py index 81ab6d9..a14e069 100644 --- a/dino-pretraining/tests/make_dummy_atlas_dataset.py +++ b/dino-pretraining/tests/make_dummy_aligned_dataset.py @@ -1,4 +1,4 @@ -"""Create a tiny raw MR-RATE-atlas NIfTI tree for integration tests.""" +"""Create a tiny raw MR-RATE coreg/atlas NIfTI tree for integration tests.""" import argparse from pathlib import Path @@ -13,13 +13,14 @@ def main() -> None: parser.add_argument("--studies", type=int, default=8) parser.add_argument("--sequences", type=int, default=3) parser.add_argument("--shape", type=int, nargs=3, default=(8, 32, 32)) + parser.add_argument("--space", choices=("coreg_space", "atlas_space"), default="coreg_space") args = parser.parse_args() d, h, w = args.shape z, y, x = np.indices((d, h, w), dtype=np.float32) affine = np.diag([0.5, 0.5, 1.0, 1.0]).astype(np.float32) for study in range(args.studies): - image_dir = Path(args.out) / "batch00" / f"dummy_{study:03d}" / "atlas_img" + image_dir = Path(args.out) / "batch00" / f"dummy_{study:03d}" / {"coreg_space": "coreg_img", "atlas_space": "atlas_img"}[args.space] image_dir.mkdir(parents=True, exist_ok=True) for sequence in range(args.sequences): center = np.array((d, h, w), dtype=np.float32) / 2 @@ -37,7 +38,7 @@ def main() -> None: nib.Nifti1Image(nifti_array, affine), image_dir / f"sequence_{sequence:02d}.nii.gz", ) - print(f"Wrote {args.studies} raw atlas studies under {args.out}") + print(f"Wrote {args.studies} raw {args.space} studies under {args.out}") if __name__ == "__main__": diff --git a/dino-pretraining/tests/make_dummy_cache.py b/dino-pretraining/tests/make_dummy_cache.py index 37ab876..faf1363 100644 --- a/dino-pretraining/tests/make_dummy_cache.py +++ b/dino-pretraining/tests/make_dummy_cache.py @@ -1,4 +1,4 @@ -"""Create a small MR-RATE-compatible atlas-space cache for GPU smoke tests.""" +"""Create a small MR-RATE-compatible aligned-space (coreg default) cache for GPU smoke tests.""" import argparse import json @@ -13,14 +13,15 @@ def main() -> None: parser.add_argument("--studies", type=int, default=8) parser.add_argument("--sequences", type=int, default=3) parser.add_argument("--shape", type=int, nargs=3, default=(64, 192, 192)) + parser.add_argument("--space", choices=("coreg_space", "atlas_space"), default="coreg_space") args = parser.parse_args() - space = Path(args.out) / "atlas_space" + space = Path(args.out) / args.space space.mkdir(parents=True, exist_ok=True) manifest = { "version": 1, "layout": "per_subject_stack", - "space": "atlas_space", + "space": args.space, "target_spacing": [1.0, 0.5, 0.5], "target_shape": list(args.shape), "posterior_shift_mm": 15.0, @@ -44,7 +45,7 @@ def main() -> None: volume += np.sin((x + sequence * 3) / 9) * 0.04 volumes.append(volume.astype(np.float16)) np.savez(space / f"dummy_{study:03d}.npz", volumes=np.stack(volumes)) - print(f"Wrote {args.studies} dummy atlas-registered studies to {space}") + print(f"Wrote {args.studies} dummy {args.space} studies to {space}") if __name__ == "__main__": diff --git a/dino-pretraining/tests/test_fora_recipe.py b/dino-pretraining/tests/test_fora_recipe.py new file mode 100644 index 0000000..3a4e4e2 --- /dev/null +++ b/dino-pretraining/tests/test_fora_recipe.py @@ -0,0 +1,174 @@ +"""Tests for the recipe fixes ported from the FORA CT DINOv3 pipeline (mr_dino/recipe.py).""" +import math + +import pytest +import torch + +from mr_dino.data import CropSpec, MRAlignedDINO3DDataset, collate_dino3d +from mr_dino.recipe import ( + CosineLinear, + GroupedAdamW, + binned_sinkhorn, + build_param_groups, + cancel_last_layer_gradients, + normalize_prototype_layer, + overlap_pairs, + position_bin_ids, + snap_second_start, + token_offset, +) +from test_mr_dino import make_dummy_cache + +PATCH = (2, 8, 8) +GLOBAL = (8, 32, 32) + + +def big_dataset(cache, cross_sequence_probability=0.0, seed=5): + return MRAlignedDINO3DDataset( + preprocessed_dir=str(cache), + crop_spec=CropSpec(global_shape=GLOBAL, local_shape=(4, 16, 16), local_crops=3), + cross_sequence_probability=cross_sequence_probability, + candidate_trials=3, + seed=seed, + global_overlap=0.25, + patch_size=PATCH, + ) + + +def test_cosine_prototypes_are_bounded_and_keep_parameter_name(): + head = torch.nn.Module() + head.last_layer = torch.nn.Linear(16, 10, bias=False) + torch.nn.init.normal_(head.last_layer.weight, std=5.0) + normalize_prototype_layer(head) + assert isinstance(head.last_layer, CosineLinear) + assert "last_layer.weight" in dict(head.named_parameters()) + x = torch.nn.functional.normalize(torch.randn(7, 16), dim=-1) + logits = head.last_layer(x) + assert logits.abs().max() <= 1.0 + 1e-5 + expected = x @ torch.nn.functional.normalize(head.last_layer.weight, dim=1).t() + torch.testing.assert_close(logits, expected) + + +def test_global_crops_overlap_quarter_and_sit_on_patch_lattice(tmp_path): + dataset = big_dataset(make_dummy_cache(tmp_path / "cache", shape=(16, 64, 64))) + moved = 0 + for i in range(len(dataset)): + sample = dataset[(i, 1)] + g0, g1 = sample["global_starts"] + token_offset(g0, g1, PATCH) # raises unless patch-aligned + for a, b, n in zip(g0, g1, GLOBAL): + assert min(a + n, b + n) - max(a, b) >= math.floor(0.25 * n) - max(PATCH) + moved += g0 != g1 + for start in sample["local_starts"]: + assert any(all(g <= s and s + l <= g + G for s, l, g, G in zip(start, (4, 16, 16), p, GLOBAL)) + for p in (g0, g1)) + assert moved > 0 # crops are no longer near-duplicates + + +def test_twin_tokens_cover_identical_voxels(tmp_path): + dataset = big_dataset(make_dummy_cache(tmp_path / "cache", shape=(16, 64, 64))) + grid = tuple(n // p for n, p in zip(GLOBAL, PATCH)) + checked = 0 + for i in range(len(dataset)): + sample = dataset[(i, 3)] + g0, g1 = sample["global_starts"] + a, b = overlap_pairs(token_offset(g0, g1, PATCH), grid) + crop0, crop1 = sample["teacher_global"][0, 0], sample["teacher_global"][1, 0] + def patch(crop, flat): + z, rem = divmod(int(flat), grid[1] * grid[2]); y, x = divmod(rem, grid[2]) + return crop[z*PATCH[0]:(z+1)*PATCH[0], y*PATCH[1]:(y+1)*PATCH[1], x*PATCH[2]:(x+1)*PATCH[2]] + for ta, tb in list(zip(a, b))[:20]: + torch.testing.assert_close(patch(crop0, ta), patch(crop1, tb)) # same sequence (p=0) + checked += 1 + assert checked > 0 + + +def test_collate_pairs_only_same_sequence_by_default(tmp_path): + dataset = big_dataset(make_dummy_cache(tmp_path / "cache", shape=(16, 64, 64)), cross_sequence_probability=1.0) + samples = [dataset[(i, 0)] for i in range(4)] + assert all(s["global_sequences"][0] != s["global_sequences"][1] for s in samples) + assert collate_dino3d(samples, patch_size=PATCH)["cross_a"].numel() == 0 + assert collate_dino3d(samples, patch_size=PATCH, cross_view_sequences="any")["cross_a"].numel() > 0 + + +def test_binned_sinkhorn_balances_each_bin_separately(): + torch.manual_seed(0) + grid, bins = (4, 4, 4), (2, 2, 2) + idx = torch.arange(64) + ids = position_bin_ids(idx, 64, grid, bins) + assert torch.bincount(ids).tolist() == [8] * 8 + logits = torch.randn(64, 16) * 3 + from mr_dino.objective import distributed_sinkhorn + q = binned_sinkhorn(distributed_sinkhorn, logits, 0.1, ids, 8) + torch.testing.assert_close(q.sum(-1), torch.ones(64)) + for b in range(8): # identical to balancing each bin on its own + rows = (ids == b).nonzero().flatten() + torch.testing.assert_close(q[rows], distributed_sinkhorn(logits[rows], 0.1)) + assert not torch.allclose(q, distributed_sinkhorn(logits, 0.1)) # and differs from one global balance + + +def test_param_groups_layerwise_decay_and_grouped_adamw(): + class Student(torch.nn.Module): + def __init__(self): + super().__init__() + self.backbone = torch.nn.Module() + self.backbone.patch_embed = torch.nn.Linear(4, 4) + self.backbone.blocks = torch.nn.ModuleList([torch.nn.Linear(4, 4) for _ in range(3)]) + self.backbone.norm = torch.nn.LayerNorm(4) + self.backbone.cls_token = torch.nn.Parameter(torch.zeros(1, 1, 4)) + self.dino_head = torch.nn.Module() + self.dino_head.last_layer = torch.nn.Linear(4, 8, bias=False) + student = Student() + groups = build_param_groups(student, layerwise_decay=0.5, patch_embed_lr_mult=0.2, n_blocks=3) + lr_of = {name: g["lr_mult"] for g in groups for name in g["names"]} + wd_of = {name: g["wd_mult"] for g in groups for name in g["names"]} + assert lr_of["backbone.patch_embed.weight"] == pytest.approx(0.5 ** 4 * 0.2) + assert lr_of["backbone.blocks.0.weight"] == pytest.approx(0.5 ** 3) + assert lr_of["backbone.blocks.2.weight"] == pytest.approx(0.5) + assert lr_of["backbone.norm.weight"] == 1.0 and lr_of["dino_head.last_layer.weight"] == 1.0 + assert wd_of["backbone.blocks.0.bias"] == 0.0 and wd_of["backbone.cls_token"] == 0.0 and wd_of["backbone.blocks.0.weight"] == 1.0 + for g in groups: + g.pop("names") + opt = GroupedAdamW(groups, lr=1.0, weight_decay=0.1) + for g in opt.param_groups: + g["lr"] = 1e-2 + student.backbone.patch_embed.weight.grad = torch.ones(4, 4) + before = student.backbone.patch_embed.weight.detach().clone() + opt.step() + assert all(g["lr"] == 1e-2 for g in opt.param_groups) # schedule untouched after step + step = (before - student.backbone.patch_embed.weight.detach()).abs().max() + assert step == pytest.approx(1e-2 * 0.5 ** 4 * 0.2, rel=0.05) # Adam step ~ lr * lr_mult + + +def test_cancel_last_layer_gradients(): + head = torch.nn.Module() + head.last_layer = torch.nn.Linear(3, 3, bias=False) + head.mlp = torch.nn.Linear(3, 3) + (head.last_layer(head.mlp(torch.randn(2, 3)))).sum().backward() + assert cancel_last_layer_gradients(head) == 1 + assert head.last_layer.weight.grad is None and head.mlp.weight.grad is not None + + +def test_cross_view_term_trains_the_backbone(tmp_path): + pytest.importorskip("dinov3") + from mr_dino.model import DinoVisionTransformer3D + from mr_dino.objective import DINO3DLearner, LossWeights + + dataset = big_dataset(make_dummy_cache(tmp_path / "cache", shape=(16, 64, 64))) + batch = collate_dino3d([dataset[(0, 1)], dataset[(1, 1)]], patch_size=PATCH) + assert batch["cross_a"].numel() > 0 + for key in ("teacher_global", "student_global", "student_local"): + batch[key] = batch[key].float() + batch["loss_weights"] = batch["sample_weights"] + torch.manual_seed(0) + backbone = DinoVisionTransformer3D(volume_size=GLOBAL, patch_size=PATCH, voxel_spacing_mm=(1.0, 0.5, 0.5), + embed_dim=96, depth=2, num_heads=3, ffn_ratio=2, n_storage_tokens=2, drop_path_rate=0) + learner = DINO3DLearner(backbone, prototypes=32, head_hidden_dim=64, dino_bottleneck_dim=32, ibot_bottleneck_dim=32, + loss_weights=LossWeights(dino=0, ibot=0, koleo=0, gram=0), position_bins=(2, 2, 2), cross_view_weight=1.0) + loss, metrics = learner(batch, teacher_temperature=0.07, step=0) + assert metrics["cross_pairs"] > 0 and torch.isfinite(metrics["cross"]) and metrics["cross"] > 0 + for key in ("dino_target_entropy", "ibot_target_entropy", "dino_kl", "ibot_kl", "ibot_prototype_usage"): + assert torch.isfinite(metrics[key]) + loss.backward() # only the cross term is weighted: it must reach the backbone blocks + grads = [p.grad for p in learner.student.backbone.blocks.parameters() if p.grad is not None] + assert grads and any(g.abs().sum() > 0 for g in grads) diff --git a/dino-pretraining/tests/test_mr_dino.py b/dino-pretraining/tests/test_mr_dino.py index 0d91cf8..324fc62 100644 --- a/dino-pretraining/tests/test_mr_dino.py +++ b/dino-pretraining/tests/test_mr_dino.py @@ -9,24 +9,24 @@ CACHE_MANIFEST_NAME, CropSpec, InfiniteStudySampler, - MRAtlasDINO3DDataset, + MRAlignedDINO3DDataset, _raw_volume, _contained_start, _intersection_box, collate_dino3d, - discover_raw_atlas, + discover_raw_aligned, previous_mr_data_module, - validate_atlas_cache, + validate_aligned_cache, ) -def make_dummy_cache(root, studies=4, shape=(8, 32, 32)): - space = root / "atlas_space" +def make_dummy_cache(root, studies=4, shape=(8, 32, 32), space_name="coreg_space"): + space = root / space_name space.mkdir(parents=True) manifest = { "version": 1, "layout": "per_subject_stack", - "space": "atlas_space", + "space": space_name, "target_spacing": [1.0, 0.5, 0.5], "target_shape": list(shape), "posterior_shift_mm": 15.0, @@ -49,7 +49,7 @@ def make_dummy_cache(root, studies=4, shape=(8, 32, 32)): def tiny_dataset(cache, seed=17): - return MRAtlasDINO3DDataset( + return MRAlignedDINO3DDataset( preprocessed_dir=str(cache), crop_spec=CropSpec( global_shape=(8, 32, 32), @@ -62,40 +62,42 @@ def tiny_dataset(cache, seed=17): ) -def test_cache_contract_rejects_non_atlas(tmp_path): +def test_cache_contract_rejects_other_space(tmp_path): cache = make_dummy_cache(tmp_path / "cache") - assert validate_atlas_cache(str(cache))["space"] == "atlas_space" - manifest_path = cache / "atlas_space" / CACHE_MANIFEST_NAME + assert validate_aligned_cache(str(cache))["space"] == "coreg_space" + with pytest.raises(FileNotFoundError): + validate_aligned_cache(str(cache), space="atlas_space") # coreg cache, atlas requested + manifest_path = cache / "coreg_space" / CACHE_MANIFEST_NAME manifest = json.loads(manifest_path.read_text()) manifest["space"] = "native_space" manifest_path.write_text(json.dumps(manifest)) - with pytest.raises(ValueError, match="expected 'atlas_space'"): - validate_atlas_cache(str(cache)) + with pytest.raises(ValueError, match="expected 'coreg_space'"): + validate_aligned_cache(str(cache)) -def test_cache_contract_rejects_coreg_space_argument(tmp_path): +def test_native_space_is_rejected(tmp_path): cache = make_dummy_cache(tmp_path / "cache") - with pytest.raises(ValueError, match="requires space='atlas_space'"): - validate_atlas_cache(str(cache), space="coreg_space") + with pytest.raises(ValueError, match="requires an aligned space"): + validate_aligned_cache(str(cache), space="native_space") + with pytest.raises(ValueError, match="requires an aligned space"): + discover_raw_aligned(str(tmp_path), space="native_space") -def test_raw_discovery_selects_atlas_img_not_coreg_img(tmp_path): - atlas_dir = tmp_path / "batch00" / "study_atlas" / "atlas_img" - atlas_dir.mkdir(parents=True) - (atlas_dir / "t1.nii.gz").touch() - (atlas_dir / "flair.nii.gz").touch() - coreg_dir = tmp_path / "batch00" / "study_coreg" / "coreg_img" - coreg_dir.mkdir(parents=True) - (coreg_dir / "t1.nii.gz").touch() - - samples = discover_raw_atlas(str(tmp_path), selected=None) - assert [sample["study_uid"] for sample in samples] == ["study_atlas"] - assert samples[0]["n_sequences"] == 2 +def test_raw_discovery_selects_requested_space_folder(tmp_path): + for name, sub, files in (("study_a", "atlas_img", 1), ("study_c", "coreg_img", 2), ("study_n", "img", 3)): + folder = tmp_path / "batch00" / name / sub + folder.mkdir(parents=True) + for i in range(files): + (folder / f"seq{i}.nii.gz").touch() + coreg = discover_raw_aligned(str(tmp_path), selected=None) # default space is coreg + assert [(s["study_uid"], s["n_sequences"]) for s in coreg] == [("study_c", 2)] + atlas = discover_raw_aligned(str(tmp_path), selected=None, space="atlas_space") + assert [(s["study_uid"], s["n_sequences"]) for s in atlas] == [("study_a", 1)] def test_raw_transform_is_exact_previous_mil_transform(tmp_path): nib = pytest.importorskip("nibabel") - image_dir = tmp_path / "batch00" / "study" / "atlas_img" + image_dir = tmp_path / "batch00" / "study" / "coreg_img" image_dir.mkdir(parents=True) x, y, z = np.indices((12, 14, 6), dtype=np.float32) array = np.sin(x / 3) + np.cos(y / 4) + z / 7 @@ -124,10 +126,11 @@ def test_aligned_cross_sequence_views_and_determinism(tmp_path): assert torch.equal(first["teacher_global"], again["teacher_global"]) assert first["teacher_global"].shape == (2, 1, 8, 32, 32) assert first["student_local"].shape == (2, 1, 4, 16, 16) - lo, hi = _intersection_box(first["global_starts"], (8, 32, 32)) - for start in first["local_starts"]: - for s, n, left, right in zip(start, (4, 16, 16), lo, hi): - assert left <= s and s + n <= right + for start in first["local_starts"]: # each local crop lies inside one of the two globals + assert any( + all(g <= s and s + n <= g + G for s, n, g, G in zip(start, (4, 16, 16), parent, (8, 32, 32))) + for parent in first["global_starts"] + ) def test_collate_masks_match_patch_grid(tmp_path): diff --git a/dino-pretraining/tests/test_probe.py b/dino-pretraining/tests/test_probe.py new file mode 100644 index 0000000..afe04e8 --- /dev/null +++ b/dino-pretraining/tests/test_probe.py @@ -0,0 +1,182 @@ +"""Validation MIL probe: checkpoint loading, token extraction, patient-level CV, metrics.""" +import csv +import json +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest +import torch + +from mr_dino.probe import ( + TEACHER_PREFIX, + TokenBags, + auroc, + average_precision, + embed_sequence, + export_teacher, + load_teacher_backbone, + main, + patient_folds, +) +from mr_dino.train_ddp import build_backbone, check_checkpoint_space + +TESTS = Path(__file__).resolve().parent + + +def write_dcp_checkpoint(path: Path, backbone, args: dict, step: int = 7) -> Path: + import torch.distributed.checkpoint as dcp + + state = {TEACHER_PREFIX + k: v.clone() for k, v in backbone.state_dict().items()} + state["model.student.backbone.cls_token"] = torch.zeros_like(backbone.cls_token) # ignored by the loader + dcp.save(state, storage_writer=dcp.FileSystemWriter(str(path)), no_dist=True) + (path / "metadata.json").write_text(json.dumps({"step": step, "stage": "pretrain", "args": args})) + (path / "COMPLETE").touch() + return path + + +def tiny_args(**extra) -> dict: + return {"arch": "tiny", "space": "coreg_space", "target_shape": [8, 32, 32], + "target_spacing": [1.0, 0.5, 0.5], "posterior_shift_mm": 0.0, **extra} + + +def test_teacher_backbone_loads_from_dcp_and_export(tmp_path): + torch.manual_seed(0) + reference = build_backbone("tiny").eval() + reference.init_weights() + ckpt = write_dcp_checkpoint(tmp_path / "step_00000007", reference, tiny_args()) + loaded, meta = load_teacher_backbone(ckpt) + assert meta["step"] == 7 and meta["args"]["arch"] == "tiny" + x = torch.randn(1, 1, 8, 32, 32) + expected = reference.forward_features(x)["x_norm_patchtokens"] + torch.testing.assert_close(loaded.forward_features(x)["x_norm_patchtokens"], expected) + exported = export_teacher(ckpt, tmp_path / "teacher.pt") + again, meta2 = load_teacher_backbone(exported) + assert meta2["step"] == 7 + torch.testing.assert_close(again.forward_features(x)["x_norm_patchtokens"], expected) + + +def test_embed_sequence_keeps_pooled_foreground_only(): + torch.manual_seed(0) + backbone = build_backbone("tiny").eval() + backbone.init_weights() + volume = torch.full((16, 64, 64), -0.4) + volume[2:10, 8:40, 8:40] = 0.3 # "head" occupies part of the volume + tokens, coords = embed_sequence(backbone, volume, tile=(8, 32, 32), pool=(2, 2, 2)) + # token grid 8x4x4, pooled 4x2x2; head covers token z 1..4, y/x 0..2 -> pooled z 0..2, y/x 0..1 + assert tokens.dtype == torch.float16 and tokens.shape[1] == backbone.embed_dim + assert {tuple(c) for c in coords.tolist()} == {(z, y, x) for z in range(3) for y in range(2) for x in range(2)} + assert torch.isfinite(tokens.float()).all() + empty, _ = embed_sequence(backbone, torch.full((8, 32, 32), -0.4), tile=(8, 32, 32), pool=(2, 2, 2)) + assert len(empty) == 0 + + +def test_patient_folds_keep_patients_together(): + uids = [f"s{i}" for i in range(40)] + patient_of = {u: f"p{i // 3}" for i, u in enumerate(uids)} + folds = patient_folds(uids, patient_of, 5, seed=1) + for p in set(patient_of.values()): + assert len({f for u, f in zip(uids, folds) if patient_of[u] == p}) == 1 + assert sorted(set(folds.tolist())) == [0, 1, 2, 3, 4] + assert (folds == patient_folds(uids, patient_of, 5, seed=1)).all() + + +def test_auroc_and_ap_match_definitions(): + rng = np.random.default_rng(0) + y = rng.integers(0, 2, 200) + s = np.round(rng.normal(size=200) + y, 1) # ties on purpose + pos, neg = s[y == 1], s[y == 0] + brute = ((pos[:, None] > neg[None]).sum() + 0.5 * (pos[:, None] == neg[None]).sum()) / (len(pos) * len(neg)) + assert auroc(y, s) == pytest.approx(brute) + assert auroc(np.zeros(5), np.arange(5.0)) is None + assert average_precision(np.array([1, 0, 1]), np.array([0.9, 0.8, 0.1])) == pytest.approx((1 + 2 / 3) / 2) + + +def write_token_shards(root: Path, n=60, dim=16, seed=0): + """Bags whose label-0 positives carry a signal token direction; two worker shards.""" + rng = np.random.default_rng(seed) + uids = [f"study{i:03d}" for i in range(n)] + labels = rng.integers(0, 2, (n, 2)) + root.mkdir(parents=True) + for rank in range(2): + mine = list(range(rank, n, 2)) + chunks, starts, counts, total = [], [], [], 0 + for i in mine: + bag = rng.normal(size=(int(rng.integers(5, 12)), dim)).astype(np.float16) + if labels[i, 0]: + bag[0, 0] += 6.0 + chunks.append(bag) + starts.append(total) + counts.append(len(bag)) + total += len(bag) + (root / f"tokens_rank{rank:04d}.f16").write_bytes(np.concatenate(chunks).tobytes()) + np.savez(root / f"index_rank{rank:04d}.npz", format="mrdino_probe_tokens_v1", dim=dim, + token_file=f"tokens_rank{rank:04d}.f16", study_uid=np.array([uids[i] for i in mine]), + start=np.array(starts), count=np.array(counts), checkpoint_step=7, checkpoint_stage="pretrain", + space="coreg_space", tile=np.array([8, 32, 32]), pool=np.array([2, 2, 2])) + with open(root.parent / "labels.csv", "w", newline="") as f: + w = csv.writer(f) + w.writerow(["study_uid", "signal", "noise"]) + w.writerows([u, *labels[i]] for i, u in enumerate(uids)) + with open(root.parent / "splits.csv", "w", newline="") as f: + w = csv.writer(f) + w.writerow(["study_uid", "patient_uid", "split"]) + w.writerows([u, f"p{i // 2}", "val"] for i, u in enumerate(uids)) + return uids + + +def test_cv_finds_signal_and_predicts_every_study_once(tmp_path): + uids = write_token_shards(tmp_path / "feat") + bags = TokenBags(tmp_path / "feat") + assert sorted(bags.uids) == uids + common = ["--features-dir", str(tmp_path / "feat"), "--labels-csv", str(tmp_path / "labels.csv"), + "--splits-csv", str(tmp_path / "splits.csv"), "--device", "cpu"] + for fold in range(3): + main(["cv", *common, "--out-dir", str(tmp_path / "out"), "--fold", str(fold), "--folds", "3", + "--epochs", "15", "--batch-size", "4", "--hidden-dim", "16", "--mlp-hidden-dim", "16"]) + main(["summarize", *common, "--out-dir", str(tmp_path / "out"), "--folds", "3"]) + result = json.loads((tmp_path / "out" / "results.json").read_text()) + assert result["studies"] == 60 and result["checkpoint_step"] == 7 + signal = next(r for r in result["per_class"] if r["label"] == "signal") + assert signal["auroc"] > 0.9 + + +def test_extract_end_to_end_on_raw_coreg_tree(tmp_path): + pytest.importorskip("nibabel") + data = tmp_path / "data" + subprocess.run([sys.executable, str(TESTS / "make_dummy_aligned_dataset.py"), "--out", str(data), + "--studies", "5", "--shape", "8", "32", "32"], check=True) + uids = sorted(p.name for p in (data / "batch00").iterdir()) + with open(tmp_path / "labels.csv", "w", newline="") as f: + w = csv.writer(f) + w.writerow(["study_uid", "a"]) + w.writerows([u, i % 2] for i, u in enumerate(uids)) + with open(tmp_path / "splits.csv", "w", newline="") as f: + w = csv.writer(f) + w.writerow(["study_uid", "patient_uid", "split"]) + w.writerows([u, u, "val" if i < 4 else "train"] for i, u in enumerate(uids)) + torch.manual_seed(0) + backbone = build_backbone("tiny").eval() + backbone.init_weights() + ckpt = write_dcp_checkpoint(tmp_path / "ckpt", backbone, tiny_args()) + args = ["extract", "--checkpoint", str(ckpt), "--data-folder", str(data), "--features-dir", str(tmp_path / "feat"), + "--labels-csv", str(tmp_path / "labels.csv"), "--splits-csv", str(tmp_path / "splits.csv"), + "--tile", "8", "32", "32", "--workers", "0", "--device", "cpu"] + main(args) + bags = TokenBags(tmp_path / "feat") + assert bags.uids == uids[:4] # val only, sorted + assert all(len(bags.tokens(i)) > 0 for i in range(len(bags))) + with np.load(tmp_path / "feat" / "index_rank0000.npz") as index: + assert set(index["seq_id"].tolist()) == {0, 1, 2} # every sequence contributes + with pytest.raises(ValueError, match="trained on 'coreg_space'"): + main([*args[:-2], "--device", "cpu", "--space", "atlas_space", "--features-dir", str(tmp_path / "f2")]) + + +def test_checkpoint_space_guard(): + check_checkpoint_space("mrrate_aligned_dinov3d_fsdp2_v1", {"space": "coreg_space"}, "coreg_space", + "mrrate_aligned_dinov3d_fsdp2_v1") + with pytest.raises(RuntimeError, match="trained on 'atlas_space'"): # pre-coreg checkpoints are atlas + check_checkpoint_space("mrrate_atlas_dinov3d_fsdp2_v1", {}, "coreg_space", "mrrate_aligned_dinov3d_fsdp2_v1") + with pytest.raises(RuntimeError, match="not an MR DINO checkpoint"): + check_checkpoint_space("something_else", {}, "coreg_space", "mrrate_aligned_dinov3d_fsdp2_v1") diff --git a/dino-pretraining/tests/test_resume_frozen_last_layer.py b/dino-pretraining/tests/test_resume_frozen_last_layer.py new file mode 100644 index 0000000..c98f472 --- /dev/null +++ b/dino-pretraining/tests/test_resume_frozen_last_layer.py @@ -0,0 +1,70 @@ +"""Resume from a checkpoint saved during freeze_last_layer (GPU smoke 941555 failed on this): +the frozen prototype layer never gets a gradient, so AdamW has no state for it and the checkpoint lacks +`optimizer.state.<...>.last_layer.*`. Only those keys may be missing; anything else must still fail loudly.""" +import pytest +import torch +import torch.distributed as dist +from torch import nn + +from dinov3.checkpointer import load_checkpoint, save_checkpoint +from mr_dino.recipe import cancel_last_layer_gradients, missing_frozen_optimizer_keys + + +@pytest.fixture(autouse=True) +def single_process_group(tmp_path): + """dinov3's DCP checkpointer needs a process group; a 1-rank gloo group is enough on CPU.""" + if not dist.is_initialized(): + dist.init_process_group("gloo", init_method=f"file://{tmp_path}/pg", rank=0, world_size=1) + yield + dist.destroy_process_group() + + +class Tiny(nn.Module): + def __init__(self): + super().__init__() + self.backbone = nn.Linear(4, 4) + self.dino_head = nn.ModuleDict({"mlp": nn.Linear(4, 4), "last_layer": nn.Linear(4, 3, bias=False)}) + + def forward(self, x): + return self.dino_head["last_layer"](self.dino_head["mlp"](self.backbone(x))) + + +def _train_step(model, opt, freeze=True, drop=None): + opt.zero_grad(set_to_none=True) + model(torch.randn(2, 4)).sum().backward() + if freeze: + cancel_last_layer_gradients(model) + if drop is not None: # simulate another parameter without optimizer state + dict(model.named_parameters())[drop].grad = None + opt.step() + + +def test_resume_during_freeze_allows_only_last_layer_state(tmp_path): + torch.manual_seed(0) + model = Tiny() + opt = torch.optim.AdamW(model.parameters(), lr=1e-2) + _train_step(model, opt) + save_checkpoint(tmp_path / "ckpt", iteration=1, model=model, optimizer=opt) + + fresh = Tiny() + fresh_opt = torch.optim.AdamW(fresh.parameters(), lr=1e-2) + missing = missing_frozen_optimizer_keys(tmp_path / "ckpt", fresh, fresh_opt) + assert missing and all(k.startswith("optimizer.state.") and ".last_layer." in k for k in missing) + with pytest.raises(BaseException, match="Missing key in checkpoint state_dict.*last_layer"): # old strict resume = GPU smoke 941555 failure + strict_model = Tiny() + load_checkpoint(tmp_path / "ckpt", model=strict_model, optimizer=torch.optim.AdamW(strict_model.parameters()), strict_loading=True) + it = load_checkpoint(tmp_path / "ckpt", model=fresh, optimizer=fresh_opt, strict_loading=not missing) + assert it == 1 + for (n, a), (_, b) in zip(model.state_dict().items(), fresh.state_dict().items()): + assert torch.equal(a, b), n + + +def test_other_missing_state_still_raises(tmp_path): + torch.manual_seed(0) + model = Tiny() + opt = torch.optim.AdamW(model.parameters(), lr=1e-2) + _train_step(model, opt, drop="dino_head.mlp.weight") + save_checkpoint(tmp_path / "ckpt", iteration=1, model=model, optimizer=opt) + fresh = Tiny() + with pytest.raises(RuntimeError, match="lacks"): + missing_frozen_optimizer_keys(tmp_path / "ckpt", fresh, torch.optim.AdamW(fresh.parameters(), lr=1e-2)) diff --git a/dino-pretraining/tests/test_stage_resume.py b/dino-pretraining/tests/test_stage_resume.py new file mode 100644 index 0000000..9dd7a7b --- /dev/null +++ b/dino-pretraining/tests/test_stage_resume.py @@ -0,0 +1,30 @@ +"""Stage transitions: an explicit RESUME (previous stage) must not override this stage's own progress.""" +from pathlib import Path + +from mr_dino.train_fsdp import resolve_checkpoint + + +def make_ckpt(root: Path, step: int, complete: bool = True) -> Path: + path = root / "checkpoints" / f"step_{step:08d}" + path.mkdir(parents=True) + if complete: + (path / "COMPLETE").touch() + (root / "checkpoints" / "latest.txt").write_text(path.name + "\n") + return path + + +def test_previous_stage_checkpoint_until_own_progress(tmp_path): + previous = make_ckpt(tmp_path / "pretrain", 50_000) + gram = tmp_path / "gram" + assert resolve_checkpoint(gram, str(previous)) == previous # first start of phase 2 + own = make_ckpt(gram, 500) + assert resolve_checkpoint(gram, str(previous)) == own # requeue continues phase 2 + assert resolve_checkpoint(gram, "latest") == own + assert resolve_checkpoint(gram, "none") is None + + +def test_incomplete_own_checkpoint_is_ignored(tmp_path): + previous = make_ckpt(tmp_path / "pretrain", 50_000) + make_ckpt(tmp_path / "gram", 500, complete=False) + assert resolve_checkpoint(tmp_path / "gram", str(previous)) == previous + assert resolve_checkpoint(tmp_path / "fresh", "latest") is None