diff --git a/README.md b/README.md index ef0653b..fc7ec57 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,10 @@ Each folder includes its own `README.md` detailing configuration, dependencies, ## Workflow Summary +For the Dong MR DINO checkpoints, full multi-sequence MIL (online/offline) and +streaming contrastive VLM are documented in +[`dino-pretraining/MIL_VLM.md`](dino-pretraining/MIL_VLM.md). + 1. **Download preprocessed data** directly from [MR-RATE on Hugging Face](https://huggingface.co/datasets/Forithmus/MR-RATE) (see `data-preprocessing/` for preprocessing details) 2. **Train the contrastive model** on (multi-sequence MRI, radiology report) pairs using `contrastive-pretraining/` 4. **Run zero-shot inference** for brain MRI pathology classification using trained model weights and `contrastive-pretraining/` diff --git a/contrastive-pretraining/README.md b/contrastive-pretraining/README.md index 368863c..ea2573b 100644 --- a/contrastive-pretraining/README.md +++ b/contrastive-pretraining/README.md @@ -7,7 +7,7 @@ Contrastive vision-language model for brain & spine MRI that aligns multi-sequen ## Architecture -Built on [FORA](https://github.com/forithmus/FORA), extending its image encoder and contrastive pretraining framework to multi-sequence brain & spine MRI with variable volumes per subject. +Built on [MR-RATE](https://github.com/forithmus/MR-RATE), extending its image encoder and contrastive pretraining framework to multi-sequence brain & spine MRI with variable volumes per subject. **Text Encoder**: [BiomedVLP-CXR-BERT-specialized](https://huggingface.co/microsoft/BiomedVLP-CXR-BERT-specialized). diff --git a/contrastive-pretraining/scripts/data.py b/contrastive-pretraining/scripts/data.py index b4e3575..5cbd08f 100644 --- a/contrastive-pretraining/scripts/data.py +++ b/contrastive-pretraining/scripts/data.py @@ -109,6 +109,13 @@ def normalize(self, data): CACHE_MANIFEST_NAME = "_manifest.json" +def is_mri_image_file(filename): + name = str(filename).lower() + auxiliary = ("_brain-mask.nii.gz", "_defacing-mask.nii.gz", "_mask.nii.gz", + "_seg.nii.gz", "_segmentation.nii.gz") + return name.endswith(".nii.gz") and not name.endswith(auxiliary) + + def list_nii_files(img_dir): """Sorted list of absolute *.nii.gz paths in img_dir (empty if it's missing).""" if not os.path.isdir(img_dir): @@ -116,7 +123,7 @@ def list_nii_files(img_dir): return sorted( os.path.join(img_dir, f) for f in os.listdir(img_dir) - if f.endswith('.nii.gz') + if is_mri_image_file(f) ) diff --git a/contrastive-pretraining/scripts/extract_features.py b/contrastive-pretraining/scripts/extract_features.py index 3a09a25..39648eb 100644 --- a/contrastive-pretraining/scripts/extract_features.py +++ b/contrastive-pretraining/scripts/extract_features.py @@ -21,7 +21,7 @@ seconds — no need to re-encode the 3D volumes every epoch. Token mode stops immediately before MR-RATE's final masked mean. It is the -required input for NeuroVFM-style Classify-Then-Aggregate MIL; pooled +required input for class-specific Classify-Then-Aggregate MIL; pooled `features_*.npy` cannot be used for MIL because the instances are gone. Usage: diff --git a/contrastive-pretraining/scripts/mil_probe.py b/contrastive-pretraining/scripts/mil_probe.py index 419a11e..3e5b652 100644 --- a/contrastive-pretraining/scripts/mil_probe.py +++ b/contrastive-pretraining/scripts/mil_probe.py @@ -1,7 +1,7 @@ """ -NeuroVFM Classify-Then-Aggregate MIL on frozen MR-RATE token features. +MRI Classify-Then-Aggregate MIL on frozen MR-RATE token features. -The head matches NeuroVFM's diagnostic architecture while replacing only +The head matches the diagnostic architecture while replacing only its visual encoder output with projected tokens from a contrastively pretrained MR-RATE checkpoint: @@ -11,7 +11,7 @@ The encoder is frozen by construction: extract_features.py runs it once under no_grad/eval and this script trains only the MIL head. Standard -nn.Linear is used instead of NeuroVFM's FlashAttention FusedDense; the +nn.Linear is used instead of the FlashAttention FusedDense; the parameterization and computation are otherwise the same. Workflow: @@ -45,7 +45,7 @@ from torch.utils.data import DataLoader, Dataset class PatchClassifier(nn.Module): - """NeuroVFM's shared MLP: Linear/GELU hidden blocks, then Linear.""" + """the shared MLP: Linear/GELU hidden blocks, then Linear.""" def __init__(self, dim: int, hidden_dims: tuple[int, ...], out_dim: int): super().__init__() @@ -62,11 +62,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class ClassifyThenAggregate(nn.Module): - """NeuroVFM's class-specific gated-attention MIL head. + """the class-specific gated-attention MIL head. Inputs are flat tokens from a batch of ragged bags plus cumulative bag boundaries. Softmax is independently normalized within each bag and - class, exactly matching NeuroVFM's segment_csr implementation. + class, exactly matching the segment_csr implementation. """ def __init__( @@ -439,7 +439,7 @@ def factor(step: int) -> float: def main() -> None: parser = argparse.ArgumentParser( - "NeuroVFM Classify-Then-Aggregate MIL on frozen MR-RATE tokens" + "MRI Classify-Then-Aggregate MIL on frozen MR-RATE tokens" ) parser.add_argument("--features_dir", type=str, required=True) parser.add_argument("--results_dir", type=str, default="./mil_probe_results") @@ -543,7 +543,7 @@ def main() -> None: use_output_bias_scale=True, ).to(device) print( - f"NeuroVFM CTA head: dim={train_data.dim}, attention_hidden={args.hidden_dim}, " + f"MRI CTA head: dim={train_data.dim}, attention_hidden={args.hidden_dim}, " f"patch_mlp=[{args.mlp_hidden_dim}], classes={n_classes}" ) print( @@ -635,7 +635,7 @@ def main() -> None: print(f"Test mean AUROC: {test_mean:.4f}") architecture = { - "name": "NeuroVFM ClassifyThenAggregate", + "name": "MRI ClassifyThenAggregate", "dim": train_data.dim, "n_classes": n_classes, "hidden_dim": args.hidden_dim, diff --git a/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_neurovfm_diagnosis_extraction/README.md b/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_mri_diagnosis_extraction/README.md similarity index 75% rename from data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_neurovfm_diagnosis_extraction/README.md rename to data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_mri_diagnosis_extraction/README.md index 1e8942e..3b8372e 100644 --- a/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_neurovfm_diagnosis_extraction/README.md +++ b/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_mri_diagnosis_extraction/README.md @@ -1,8 +1,8 @@ -# NeuroVFM MRI diagnosis extraction +# MRI MRI diagnosis extraction This is the report-to-label pipeline used for the current MR-RATE clinical evaluation target. It classifies each report against the 74 expert-defined -NeuroVFM MRI diagnoses and returns a binary label plus a short rationale for +MRI MRI diagnoses and returns a binary label plus a short rationale for every diagnosis. The checked-in prompt, diagnosis order, guidance, JSON parser, and inference @@ -13,15 +13,15 @@ the repository only to reproduce the historical 37-pathology labels. Run one or more Slurm ranks with: ```bash -srun --gpu-bind=none python extract_neurovfm_dx_gemma.py \ +srun --gpu-bind=none python extract_mri_dx_gemma.py \ --reports_dir /path/to/reports \ - --diagnoses_json data/neurovfm_mri_diagnoses.json \ + --diagnoses_json data/mri_diagnoses_74.json \ --output_dir /path/to/shards python merge_labels.py \ --input_dir /path/to/shards \ - --output /path/to/neurovfm74_labels.csv \ - --rationales /path/to/neurovfm74_rationales.json + --output /path/to/mri74_labels.csv \ + --rationales /path/to/mri74_rationales.json ``` Input files are named `batchNN_reports.csv` and require `study_uid`. The diff --git a/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_mri_diagnosis_extraction/data/mri_diagnoses_74.json b/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_mri_diagnosis_extraction/data/mri_diagnoses_74.json new file mode 100644 index 0000000..3328216 --- /dev/null +++ b/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_mri_diagnosis_extraction/data/mri_diagnoses_74.json @@ -0,0 +1,301 @@ +{ + "_comment": "Ordered 74-class MRI diagnosis schema. Preserve diagnosis keys, order, and clinical guidance.", + "diagnoses": [ + { + "key": "subdural_hematoma", + "guidance": "" + }, + { + "key": "epidural_hematoma", + "guidance": "" + }, + { + "key": "brain_contusion", + "guidance": "" + }, + { + "key": "diffuse_axonal_injury", + "guidance": "" + }, + { + "key": "cerebral_aneurysm", + "guidance": "" + }, + { + "key": "arteriovenous_malformation", + "guidance": "" + }, + { + "key": "cavernous_malformation_cavernoma", + "guidance": "\"yes\" if radiologist confident for cavernoma; \"no\" if possible microbleed" + }, + { + "key": "venous_sinus_thrombosis", + "guidance": "" + }, + { + "key": "acute_ischemic_stroke", + "guidance": "\"yes\" for acute ischemia/infarct with diffusion restriction" + }, + { + "key": "chronic_ischemic_stroke", + "guidance": "\"yes\" for chronic or evolving ischemia/infarct without diffusion restriction" + }, + { + "key": "moya_moya_disease", + "guidance": "" + }, + { + "key": "small_vessel_ischemic_disease", + "guidance": "also known as microvascular ischemic disease" + }, + { + "key": "lacunar_stroke", + "guidance": "" + }, + { + "key": "intracranial_hemorrhage", + "guidance": "\"yes\" for any intracranial hemorrhage, excluding microhemorrhage" + }, + { + "key": "intraparenchymal_hemorrhage", + "guidance": "\"yes\" for hemorrhages within the brain, usually spontaneous and due to hypertension" + }, + { + "key": "intraventricular_hemorrhage", + "guidance": "" + }, + { + "key": "aneurysmal_subarachnoid_hemorrhage", + "guidance": "\"yes\" if due to ruptured aneurysm or other vascular cause" + }, + { + "key": "traumatic_subarachnoid_hemorrhage", + "guidance": "\"yes\" if due to head trauma" + }, + { + "key": "multiple_sclerosis", + "guidance": "" + }, + { + "key": "neuromyelitis_optica", + "guidance": "" + }, + { + "key": "neurosarcoidosis", + "guidance": "" + }, + { + "key": "brain_abscess", + "guidance": "\"yes\" if radiologist is confident for brain abscesses or subdural empyemas" + }, + { + "key": "viral_encephalitis", + "guidance": "most commonly herpes encephalitis involving the temporal and frontal lobes" + }, + { + "key": "arachnoid_cyst", + "guidance": "" + }, + { + "key": "pineal_cyst", + "guidance": "" + }, + { + "key": "epidermoid_cyst", + "guidance": "\"yes\" for intracranial epidermoid cyst only" + }, + { + "key": "colloid_cyst", + "guidance": "" + }, + { + "key": "brain_tumor", + "guidance": "\"yes\" for any intracranial mass, growth, or neoplasm" + }, + { + "key": "orbital_tumor", + "guidance": "\"yes\" for any orbital mass, growth, or neoplasm" + }, + { + "key": "head_neck_tumor", + "guidance": "\"yes\" for tumor/cancer in the oral cavity, throat, larynx, nasal cavity, salivary, and thyroid gland" + }, + { + "key": "spine_tumor", + "guidance": "\"yes\" for any spine mass, growth, or neoplasm" + }, + { + "key": "nerve_tumor", + "guidance": "\"yes\" for nerve, nerve root, or nerve sheath tumors" + }, + { + "key": "spinal_degenerative_changes", + "guidance": "" + }, + { + "key": "glioma", + "guidance": "\"yes\" if mostly likely primary glial tumor" + }, + { + "key": "high_grade_glioma", + "guidance": "\"yes\" if MRI features for high-grade glioma" + }, + { + "key": "low_grade_glioma", + "guidance": "\"yes\" if MRI features for low-grade glioma" + }, + { + "key": "brain_metastasis", + "guidance": "" + }, + { + "key": "meningioma", + "guidance": "" + }, + { + "key": "schwannoma", + "guidance": "" + }, + { + "key": "lymphoma", + "guidance": "" + }, + { + "key": "pineal_tumor", + "guidance": "\"yes\" for any tumor in the pineal region" + }, + { + "key": "intraventricular_tumor", + "guidance": "" + }, + { + "key": "pediatric_posterior_fossa_tumor", + "guidance": "diagnosis includes, but not limited to, medulloblastoma, ependymoma, pilocytic astrocytoma" + }, + { + "key": "brainstem_glioma", + "guidance": "diagnosis also called diffuse midline gliomas" + }, + { + "key": "pituitary_tumor", + "guidance": "" + }, + { + "key": "rathkes_cleft_cyst", + "guidance": "" + }, + { + "key": "craniopharyngioma", + "guidance": "" + }, + { + "key": "chordoma_chondrosarcoma", + "guidance": "" + }, + { + "key": "chiari_malformation", + "guidance": "\"yes\" for chiari I/II malformations" + }, + { + "key": "cerebral_atrophy", + "guidance": "" + }, + { + "key": "brain_herniation", + "guidance": "" + }, + { + "key": "brain_mass_effect", + "guidance": "" + }, + { + "key": "brain_midline_shift", + "guidance": "" + }, + { + "key": "cephaloceles", + "guidance": "" + }, + { + "key": "encephalomalacia_gliosis", + "guidance": "" + }, + { + "key": "cerebral_edema", + "guidance": "" + }, + { + "key": "cavum_septum_pellucidum", + "guidance": "" + }, + { + "key": "dandy_walker_malformation", + "guidance": "" + }, + { + "key": "mega_cisterna_magna", + "guidance": "" + }, + { + "key": "heterotopia", + "guidance": "" + }, + { + "key": "pachygyria_lissencephaly", + "guidance": "" + }, + { + "key": "schizencephaly", + "guidance": "" + }, + { + "key": "dysgenesis_corpus_callosum", + "guidance": "" + }, + { + "key": "ventriculomegaly", + "guidance": "\"yes\" for any abnormally enlarged ventricles" + }, + { + "key": "hydrocephalus_ex_vacuo", + "guidance": "\"yes\" for ventriculomegaly due to brain atrophy, normal pressure hydrocephalus, or ex vacuo dilation" + }, + { + "key": "obstructive_hydrocephalus", + "guidance": "\"yes\" for ventriculomegaly due to mass lesion or aqueductal stenosis" + }, + { + "key": "subdural_hygroma_effusions", + "guidance": "" + }, + { + "key": "intracranial_hypotension", + "guidance": "" + }, + { + "key": "craniotomy_craniectomy", + "guidance": "" + }, + { + "key": "tumor_resection_cavity", + "guidance": "" + }, + { + "key": "postsurgical_changes", + "guidance": "" + }, + { + "key": "pneumocephalus", + "guidance": "" + }, + { + "key": "catheter", + "guidance": "\"yes\" for any intracranial catheter" + }, + { + "key": "spine_syrinx", + "guidance": "" + } + ] +} diff --git a/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_neurovfm_diagnosis_extraction/extract_neurovfm_dx_gemma.py b/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_mri_diagnosis_extraction/extract_mri_dx_gemma.py similarity index 96% rename from data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_neurovfm_diagnosis_extraction/extract_neurovfm_dx_gemma.py rename to data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_mri_diagnosis_extraction/extract_mri_dx_gemma.py index cdeaf0d..2436406 100644 --- a/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_neurovfm_diagnosis_extraction/extract_neurovfm_dx_gemma.py +++ b/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_mri_diagnosis_extraction/extract_mri_dx_gemma.py @@ -1,17 +1,17 @@ # ======================================================================= -# NeuroVFM MRI diagnosis extraction on MR-RATE reports, via Gemma (vLLM). +# MRI MRI diagnosis extraction on MR-RATE reports, via Gemma (vLLM). # -# Faithful reproduction of the NeuroVFM "MRI diagnosis extraction prompt" +# Faithful reproduction of the MRI "MRI diagnosis extraction prompt" # (Supplementary Data Figure 6): an LLM-based annotation pipeline that # converts free-text radiology reports into structured labels for the 74 # expert-defined MRI diagnoses. For EACH diagnosis the model returns a # present/absent label ("yes"/"no") plus a one-sentence rationale. # -# Differences from the original NeuroVFM pipeline: +# Differences from the original MRI pipeline: # - Runs a local open-weights Gemma model through vLLM instead of the # hosted GPT-4.1-mini (`--model_name`, default google/gemma-4-31B-it). # - Gemma has NO system role, so the instruction block + report are sent -# as a single user turn (which is how the NeuroVFM prompt is structured +# as a single user turn (which is how the MRI prompt is structured # anyway: one instruction preamble followed by the report). # # Input: MR-RATE structured report CSVs (batch{NN}_reports.csv) with @@ -67,12 +67,12 @@ # ============================================================ -# Prompt (verbatim NeuroVFM MRI diagnosis extraction prompt) +# Prompt (verbatim MRI MRI diagnosis extraction prompt) # ============================================================ # The instruction preamble. The diagnoses list is appended where the caption # reads "", then the report is appended after. -NEUROVFM_PROMPT = """\ +MRI_DIAGNOSIS_PROMPT = """\ You, ChatGPT, are an expert neuroradiologist specializing in the interpretation of brain and spine MRI imaging reports. You will classify a radiology report according to the predefined list diagnoses below. This is a multilabel classification task, meaning a patient may have any number of these diagnoses (including none). Your job is to: @@ -272,7 +272,7 @@ def main(): from vllm import LLM, SamplingParams ap = argparse.ArgumentParser( - description="Extract the 74 NeuroVFM MRI diagnoses from MR-RATE " + description="Extract the 74 MRI MRI diagnoses from MR-RATE " "reports with a local Gemma model via vLLM." ) ap.add_argument("--reports_dir", required=True, @@ -280,8 +280,8 @@ def main(): "(e.g. MR-RATE-validation/reports)") ap.add_argument("--diagnoses_json", default=os.path.join(os.path.dirname(__file__), - "neurovfm_mri_diagnoses.json"), - help="JSON with the 74 NeuroVFM diagnoses + guidance") + "mri_diagnoses_74.json"), + help="JSON with the 74 MRI diagnoses + guidance") ap.add_argument("--output_dir", required=True, help="Where to write labels_rank_{rank}.json") ap.add_argument("--model_name", default="google/gemma-4-31B-it", @@ -332,7 +332,7 @@ def main(): diagnoses = json.load(f)["diagnoses"] diagnosis_keys = [d["key"] for d in diagnoses] diagnoses_block = build_diagnoses_block(diagnoses) - prompt_preamble = NEUROVFM_PROMPT.format(diagnoses_block=diagnoses_block) + prompt_preamble = MRI_DIAGNOSIS_PROMPT.format(diagnoses_block=diagnoses_block) n_dx = len(diagnosis_keys) print(f"[Rank {rank}] Loaded {n_dx} diagnoses", flush=True) @@ -410,7 +410,7 @@ def save_output(): out = { "metadata": { "model": args.model_name, - "prompt": "neurovfm_mri_diagnosis_extraction", + "prompt": "mri_mri_diagnosis_extraction", "seed": args.seed, "n_diagnoses": n_dx, "diagnosis_keys": diagnosis_keys, diff --git a/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_neurovfm_diagnosis_extraction/merge_labels.py b/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_mri_diagnosis_extraction/merge_labels.py similarity index 91% rename from data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_neurovfm_diagnosis_extraction/merge_labels.py rename to data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_mri_diagnosis_extraction/merge_labels.py index 2e2b4cd..e844711 100644 --- a/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_neurovfm_diagnosis_extraction/merge_labels.py +++ b/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_mri_diagnosis_extraction/merge_labels.py @@ -1,13 +1,13 @@ #!/usr/bin/env python3 """ -Merge per-rank labels_rank_*.json (from extract_neurovfm_dx_gemma.py) into a -single wide CSV: study_uid + one 0/1 column per NeuroVFM diagnosis. +Merge per-rank labels_rank_*.json (from extract_mri_dx_gemma.py) into a +single wide CSV: study_uid + one 0/1 column per MRI diagnosis. Also writes an optional *_rationales.json keyed by study_uid if --rationales is given. Usage: - python merge_labels.py --input_dir /path/to/output --output neurovfm_labels.csv + python merge_labels.py --input_dir /path/to/output --output mri_labels.csv """ import argparse import csv diff --git a/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_neurovfm_diagnosis_extraction/data/neurovfm_mri_diagnoses.json b/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_neurovfm_diagnosis_extraction/data/neurovfm_mri_diagnoses.json deleted file mode 100644 index fe3ee9d..0000000 --- a/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_neurovfm_diagnosis_extraction/data/neurovfm_mri_diagnoses.json +++ /dev/null @@ -1,79 +0,0 @@ -{ - "_comment": "The 74 expert-defined MRI diagnoses from NeuroVFM Supplementary Data Figure 6 (MRI diagnosis extraction prompt). Order and key names are verbatim from the caption. 'guidance' holds the parenthetical clarification given in the caption (empty string if none). Injected into the extraction prompt where the caption reads ''.", - "diagnoses": [ - {"key": "subdural_hematoma", "guidance": ""}, - {"key": "epidural_hematoma", "guidance": ""}, - {"key": "brain_contusion", "guidance": ""}, - {"key": "diffuse_axonal_injury", "guidance": ""}, - {"key": "cerebral_aneurysm", "guidance": ""}, - {"key": "arteriovenous_malformation", "guidance": ""}, - {"key": "cavernous_malformation_cavernoma", "guidance": "\"yes\" if radiologist confident for cavernoma; \"no\" if possible microbleed"}, - {"key": "venous_sinus_thrombosis", "guidance": ""}, - {"key": "acute_ischemic_stroke", "guidance": "\"yes\" for acute ischemia/infarct with diffusion restriction"}, - {"key": "chronic_ischemic_stroke", "guidance": "\"yes\" for chronic or evolving ischemia/infarct without diffusion restriction"}, - {"key": "moya_moya_disease", "guidance": ""}, - {"key": "small_vessel_ischemic_disease", "guidance": "also known as microvascular ischemic disease"}, - {"key": "lacunar_stroke", "guidance": ""}, - {"key": "intracranial_hemorrhage", "guidance": "\"yes\" for any intracranial hemorrhage, excluding microhemorrhage"}, - {"key": "intraparenchymal_hemorrhage", "guidance": "\"yes\" for hemorrhages within the brain, usually spontaneous and due to hypertension"}, - {"key": "intraventricular_hemorrhage", "guidance": ""}, - {"key": "aneurysmal_subarachnoid_hemorrhage", "guidance": "\"yes\" if due to ruptured aneurysm or other vascular cause"}, - {"key": "traumatic_subarachnoid_hemorrhage", "guidance": "\"yes\" if due to head trauma"}, - {"key": "multiple_sclerosis", "guidance": ""}, - {"key": "neuromyelitis_optica", "guidance": ""}, - {"key": "neurosarcoidosis", "guidance": ""}, - {"key": "brain_abscess", "guidance": "\"yes\" if radiologist is confident for brain abscesses or subdural empyemas"}, - {"key": "viral_encephalitis", "guidance": "most commonly herpes encephalitis involving the temporal and frontal lobes"}, - {"key": "arachnoid_cyst", "guidance": ""}, - {"key": "pineal_cyst", "guidance": ""}, - {"key": "epidermoid_cyst", "guidance": "\"yes\" for intracranial epidermoid cyst only"}, - {"key": "colloid_cyst", "guidance": ""}, - {"key": "brain_tumor", "guidance": "\"yes\" for any intracranial mass, growth, or neoplasm"}, - {"key": "orbital_tumor", "guidance": "\"yes\" for any orbital mass, growth, or neoplasm"}, - {"key": "head_neck_tumor", "guidance": "\"yes\" for tumor/cancer in the oral cavity, throat, larynx, nasal cavity, salivary, and thyroid gland"}, - {"key": "spine_tumor", "guidance": "\"yes\" for any spine mass, growth, or neoplasm"}, - {"key": "nerve_tumor", "guidance": "\"yes\" for nerve, nerve root, or nerve sheath tumors"}, - {"key": "spinal_degenerative_changes", "guidance": ""}, - {"key": "glioma", "guidance": "\"yes\" if mostly likely primary glial tumor"}, - {"key": "high_grade_glioma", "guidance": "\"yes\" if MRI features for high-grade glioma"}, - {"key": "low_grade_glioma", "guidance": "\"yes\" if MRI features for low-grade glioma"}, - {"key": "brain_metastasis", "guidance": ""}, - {"key": "meningioma", "guidance": ""}, - {"key": "schwannoma", "guidance": ""}, - {"key": "lymphoma", "guidance": ""}, - {"key": "pineal_tumor", "guidance": "\"yes\" for any tumor in the pineal region"}, - {"key": "intraventricular_tumor", "guidance": ""}, - {"key": "pediatric_posterior_fossa_tumor", "guidance": "diagnosis includes, but not limited to, medulloblastoma, ependymoma, pilocytic astrocytoma"}, - {"key": "brainstem_glioma", "guidance": "diagnosis also called diffuse midline gliomas"}, - {"key": "pituitary_tumor", "guidance": ""}, - {"key": "rathkes_cleft_cyst", "guidance": ""}, - {"key": "craniopharyngioma", "guidance": ""}, - {"key": "chordoma_chondrosarcoma", "guidance": ""}, - {"key": "chiari_malformation", "guidance": "\"yes\" for chiari I/II malformations"}, - {"key": "cerebral_atrophy", "guidance": ""}, - {"key": "brain_herniation", "guidance": ""}, - {"key": "brain_mass_effect", "guidance": ""}, - {"key": "brain_midline_shift", "guidance": ""}, - {"key": "cephaloceles", "guidance": ""}, - {"key": "encephalomalacia_gliosis", "guidance": ""}, - {"key": "cerebral_edema", "guidance": ""}, - {"key": "cavum_septum_pellucidum", "guidance": ""}, - {"key": "dandy_walker_malformation", "guidance": ""}, - {"key": "mega_cisterna_magna", "guidance": ""}, - {"key": "heterotopia", "guidance": ""}, - {"key": "pachygyria_lissencephaly", "guidance": ""}, - {"key": "schizencephaly", "guidance": ""}, - {"key": "dysgenesis_corpus_callosum", "guidance": ""}, - {"key": "ventriculomegaly", "guidance": "\"yes\" for any abnormally enlarged ventricles"}, - {"key": "hydrocephalus_ex_vacuo", "guidance": "\"yes\" for ventriculomegaly due to brain atrophy, normal pressure hydrocephalus, or ex vacuo dilation"}, - {"key": "obstructive_hydrocephalus", "guidance": "\"yes\" for ventriculomegaly due to mass lesion or aqueductal stenosis"}, - {"key": "subdural_hygroma_effusions", "guidance": ""}, - {"key": "intracranial_hypotension", "guidance": ""}, - {"key": "craniotomy_craniectomy", "guidance": ""}, - {"key": "tumor_resection_cavity", "guidance": ""}, - {"key": "postsurgical_changes", "guidance": ""}, - {"key": "pneumocephalus", "guidance": ""}, - {"key": "catheter", "guidance": "\"yes\" for any intracranial catheter"}, - {"key": "spine_syrinx", "guidance": ""} - ] -} diff --git a/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/README.md b/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/README.md index 2d1ab64..5507ddd 100644 --- a/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/README.md +++ b/data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/README.md @@ -259,15 +259,15 @@ python 06_pathology_classification/merge_labels.py \ --output /path/to/labels.csv ``` -### 07_neurovfm_diagnosis_extraction +### 07_mri_diagnosis_extraction -This is the current clinical-evaluation labeler. It uses the NeuroVFM MRI +This is the current clinical-evaluation labeler. It uses the MRI MRI diagnosis extraction prompt and the bundled 74-diagnosis schema, including the per-diagnosis guidance used for the completed MR-RATE re-extraction. It classifies the explicitly marked FINDINGS and IMPRESSION sections, emits a binary label and rationale for every key, retries malformed JSON once, and checkpoints atomically after every batch. See -`07_neurovfm_diagnosis_extraction/README.md` for the run contract. +`07_mri_diagnosis_extraction/README.md` for the run contract. The 37-label pipeline in step 06 is retained for historical reproducibility; report-generation clinical metrics use step 07 and reject the step-06 schema. diff --git a/dino-pretraining/MIL_VLM.md b/dino-pretraining/MIL_VLM.md new file mode 100644 index 0000000..8e3be9c --- /dev/null +++ b/dino-pretraining/MIL_VLM.md @@ -0,0 +1,188 @@ +# MR DINO → MIL and VLM + +Multi-sequence MRI classification and streamed image–text contrastive training. +The image path trains the final DINO block and a visual projector; the text path +uses CXR-BERT. Two-pass pooling preserves gradients across the full study bag. + +Training is intended to run at NVIDIA. No full MRI training was started on the +development cluster. + +## Inputs and geometry + +- Each study includes every eligible MRI sequence, encoded separately. All + sequence tokens enter one study bag; sequence logits are never averaged. +- Auxiliary brain/defacing/segmentation masks are excluded from MRI inputs. +- Preprocessing uses the existing MR-RATE loader: canonical RAS, resampling, + z-score normalization, posterior shift, and crop/padding. Encoder inputs are + unchanged. Pooling support is derived before normalization, so normalized air + and zero crop padding are not mistaken for anatomy. +- The checkpoint determines its model source, architecture, prefix tokens, + physical RoPE, spacing, and training tile size. There is no CT block-24 or + 1024-dimensional assumption. Non-overlapping tiles cover the complete volume. +- The two H+ releases explicitly record `atlas_space`. The legacy 7B highres + release states atlas space in its README but omits full run arguments. Its + preprocessing fallback is the documented default at its pinned MR-RATE commit + `43c05c61bf98d00c4e825a27a0fca3c742754d78`. Confirm original run overrides; + an optional `preprocessing.json` beside its weights overrides those defaults. + +Production MIL uses the final **74-class MRI diagnosis labels**, not the +14-group development labels. Set `LABELS` to the final 74-class CSV, with +`study_uid` plus the 74 diagnosis columns. The bundled ordered schema is +`data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_mri_diagnosis_extraction/data/mri_diagnoses_74.json`. +The loader rejects missing or extra classes and reorders CSV columns to this +canonical order. The checkpoint records that class order. Values below zero or +NaN are unknown and masked; 0/1 are supervised. + +Train/validation/test patients must be disjoint. Test data is not used in training +or normalization. MRI anatomy routes may be supplied explicitly; CT organ +definitions are not imported. The default route allows all supported MRI tokens. + +## Environment + +Use Python 3.10+, the official DINOv3 checkout pinned by the release, and a +compatible CUDA PyTorch environment: + +```bash +python -m pip install -r dino-pretraining/requirements.txt +export PYTHONPATH="$PWD/dino-pretraining:$DINOV3_ROOT:$PYTHONPATH" +git lfs pull --include=contrastive-pretraining/data/findings_sentences.jsonl +``` + +`DINOV3_ROOT` must point to official DINOv3 commit +`6876159a11b4df116f30f667f8c9888617df0751`. Training fails clearly on a sentence +file that is still a Git LFS pointer. Stage large models once per node to a local +cache with a completion marker, rather than copying them per GPU. + +The examples below use eight ranks on one node. For multiple nodes, use the +site's `torchrun --nnodes ... --node-rank ... --master-addr ... --master-port ...` +settings. All ranks must see the same complete image/cache cohort. First perform +a two-node smoke in NVIDIA's actual runtime before a large launch; development +validation does not establish NVIDIA multi-node throughput or NCCL behavior. + +## Offline prefix cache + +Set `RELEASE`, `DATA`, `SPLITS`, `LABELS`, and `CACHE` to NVIDIA paths. Prefix +tokens, including each tile's prefix/context tokens, are stored as FP16. The +trainable last block is replayed during training. No token or sequence cap is +applied. This cache can be reused by MIL and VLM from the same DINO release. +It is large; use online mode if storage is constrained. + +```bash +torchrun --standalone --nproc-per-node=8 -m mr_dino.mil cache \ + --release "$RELEASE" --data-folder "$DATA" --splits-csv "$SPLITS" \ + --split train --cache-dir "$CACHE" +torchrun --standalone --nproc-per-node=8 -m mr_dino.mil cache \ + --release "$RELEASE" --data-folder "$DATA" --splits-csv "$SPLITS" \ + --split val --cache-dir "$CACHE" +``` + +Edit all paths in `configs/vlm_mri.json`. Compute the shared normalization +using **training studies only**, with the DINO-initialized final block: + +```bash +torchrun --standalone --nproc-per-node=8 -m mr_dino.vlm normalize \ + --config dino-pretraining/configs/vlm_mri.json \ + --output /path/to/mri_normalization.pt +``` + +For a fully online workflow set `mode` to `online`; normalization then reads raw +images instead of a prefix cache. + +## MIL + +The trainable final MRI block feeds normalized features and global-grid Fourier +positions to a gated classify-then-aggregate scalar head. The class softmax +spans all tokens of all sequences. Two passes give exact gradients without +retaining all transformer graphs. Localization is positive +`attention × token logit × output scale` evidence, excluding the output bias. + +```bash +torchrun --standalone --nproc-per-node=8 -m mr_dino.mil train \ + --mode offline --release "$RELEASE" --cache-dir "$CACHE" \ + --normalization /path/to/mri_normalization.pt \ + --labels-csv "$LABELS" --splits-csv "$SPLITS" \ + --head-lr 0.002 --block-lr 0.00001 --epochs 10 --out runs/mri_mil +``` + +For online MIL replace `--mode offline --cache-dir "$CACHE"` with +`--mode online --data-folder "$DATA"`. The prefix remains frozen; the final +block and MIL head train. Online and offline use the same support and geometry. +Incomplete labelled train/validation cohorts fail by default; the explicit +`--allow-incomplete-data` option is for implementation smoke tests only. + +MIL writes validation probabilities, per-class/macro AUROC and complete +optimizer/model checkpoints at epoch boundaries. `--resume epoch_001.pt` +continues from the next epoch; it does not claim mid-epoch replay. Learning rates +are explicit constant settings in this CLI. The example values are starting +settings, not a claim of optimized MRI performance. + +## VLM + +```bash +torchrun --standalone --nproc-per-node=8 -m mr_dino.vlm train \ + --config dino-pretraining/configs/vlm_mri.json --output runs/mri_vlm +``` + +Both `mode: offline` and `mode: online` implement the same objective: + +1. DINO-initialized trainable final block, fixed train-only mean/RMS, + global-grid Fourier features, and an MRI sequence-kind embedding. +2. CXR-BERT CLS plus projection into the shared space. +3. Exact streamed VL-CABS pooling over all sequence tokens, with learned pooling + and contrastive temperatures. +4. the symmetric multi-positive InfoNCE: a study's sentences are all positives; + its own other positive sentences do not become image-to-text negatives. + +The supplied learning-rate groups, cosine/warmup schedule, 34-sentence sampling, +and 128-token sentence encoding are explicit in the example configuration. These limits are +on contrastive finding statements, not generated report lengths. All visual +tokens are retained. The VLM starts its last block from DINO, independently of +MIL. MRI sequence kinds replace CT window-kind embeddings. Unknown sequence +names use an explicit `other` kind. + +Training covers eligible train studies without repeating padding examples; +inactive tail ranks are masked from positives and negatives. A singleton tail +contains no contrastive negative and has zero contrastive loss. Validation is +run each epoch on the validation cohort. Checkpoints include the optimizer, +next batch, temperatures, and per-rank RNG. `latest.pt` is resumed automatically +only within the explicitly supplied output directory. Recipe/world changes fail +instead of silently changing the resumed experiment. `SIGUSR1` requests a +checkpoint and continues training; `SIGTERM` requests a checkpoint and exits. +The trainer contains no Slurm cancellation, allocation-release, or 24-hour kill +timer. + +## Development validation and previews + +- Dense and two-pass MIL outputs/gradients agree, including gradients through + trainable features and soft routes. +- The MRI adapter matches the backbone's final block with dynamic prefix/grid. +- MR-RATE streamed VL-CABS matches dense outputs and every gradient, including both + temperatures and text/visual paths. +- A two-rank Gloo step matches the dense global contrastive objective. +- A synthetic full CLI test exercised train-only normalization, training, + validation, inactive tail ranks, and interruption/resume. Resumed final block, + projector, BERT and temperatures were bitwise identical to the uninterrupted + control. +- A real five-sequence MRI study exercised final-block/MIL backward and a weight + update. This was an implementation check, not a clinical training result. + +Run unit checks with: + +```bash +python -m pytest -q dino-pretraining/tests/test_mil_stream.py \ + dino-pretraining/tests/test_vlm_stream.py +``` + +`python -m mr_dino.preview --release ... --manifest ... --out ...` produces +whole-volume three-component PCA, fitted independently per sequence after tiles +are assembled, with physical voxel aspect ratios. The example manifest contains +`study_uid` and every `nifti_files` path. Prototype heads are absent from the +downloaded inference releases; the preview does not invent prototype maps. + +Full synthetic two-rank normalization/train/validation/resume smoke (CPU, no model download): + +```bash +bash dino-pretraining/scripts/smoke_mil_vlm.sh /tmp/mrrate_vlm_smoke +``` + +Use a fresh smoke directory. It compares resumed weights bitwise with an uninterrupted control. diff --git a/dino-pretraining/README.md b/dino-pretraining/README.md index 60600c4..27a9bf5 100644 --- a/dino-pretraining/README.md +++ b/dino-pretraining/README.md @@ -1,6 +1,6 @@ # Co-registered 3-D DINOv3 for MR-RATE -This module adapts FORA's CT DINOv3 training strategy to MR-RATE's +This module trains a native 3-D DINOv3 encoder on **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. @@ -69,45 +69,16 @@ 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) +## Training recipe -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. +The default recipe uses cosine prototype heads, overlapping global views, +patch-aligned cross-crop supervision, spatially balanced patch targets, EMA +teacher updates, and per-objective entropy/KL diagnostics. The MRI default is +ViT-H+ with 16,384 prototypes per head and 2048/256 hidden/bottleneck widths. +Cross-sequence global pairs default to 0.25. Patch-level twin-token targets +use the same sequence, preserving contrast-specific findings. -| 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. +Production launchers and checkpoint probes are described below. ## Optional input-volume cache @@ -169,10 +140,10 @@ 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 +- **Why not 7B:** 7B requires a larger training memory budget, 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. + this implementation, with lower capacity than H+. - **Learning rate:** it was tuned on ViT-L. Watch the target metrics in the first 2–3k steps (see "Healthy early training" above). @@ -192,7 +163,7 @@ heads are 2048 → 256 with 16,384 cosine prototypes each. | 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`, +Phases 2 and 3 use MR 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. @@ -236,7 +207,7 @@ STAGE=gram RESUME=$R/pretrain/checkpoints/step_00050000 OUTPUT=$R/gram sbatch sc nohup scripts/probe_watch.sh $R/gram 5000 10000 & # 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) +# otherwise RESUME=$R/pretrain/checkpoints/step_00050000 (MRI: 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 & ``` @@ -263,7 +234,7 @@ nohup scripts/probe_watch.sh $R/highres 2500 5000 & ## Validation probe (checkpoint scoring) -This is the FORA CT-DINO protocol: frozen features, MIL on the validation split only, and +This is the MR 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 @@ -275,7 +246,7 @@ patient-level cross-validation. Train and test are never touched. 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 +3. **CV** (`python -m mr_dino.probe cv --fold k`): trains MR-RATE's MRI `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`, diff --git a/dino-pretraining/configs/vlm_mri.json b/dino-pretraining/configs/vlm_mri.json new file mode 100644 index 0000000..f2bc22b --- /dev/null +++ b/dino-pretraining/configs/vlm_mri.json @@ -0,0 +1,30 @@ +{ + "mode": "offline", + "release": "/path/to/mrdino3d_hplus_atlas_infer_highres_step5000", + "cache_dir": "/path/to/mri_prefix_cache", + "data_folder": "/path/to/MR-RATE-atlas/mri", + "splits_csv": "/path/to/splits.csv", + "sentences": "/path/to/findings_sentences.jsonl", + "normalization": "/path/to/mri_normalization.pt", + "allow_incomplete_data": false, + "training": { + "epochs": 10, + "seed": 1, + "warmup_steps": 200, + "position_freqs": 4, + "max_sentences": 34, + "max_tokens": 128, + "save_every": 200, + "model": { + "dim": 512, + "text_model": "/path/to/BiomedVLP-CXR-BERT-specialized", + "lr_block": 0.00001, + "lr_proj": 0.00005, + "lr_text": 0.000005, + "lr_temp": 0.0005, + "wd": 0.05, + "pool_temperature_init": 14.285714285714285, + "position_scale": 1.0 + } + } +} diff --git a/dino-pretraining/mr_dino/data.py b/dino-pretraining/mr_dino/data.py index f018d3c..bb91baa 100644 --- a/dino-pretraining/mr_dino/data.py +++ b/dino-pretraining/mr_dino/data.py @@ -376,6 +376,26 @@ def _raw_volume( return torch.from_numpy(np.ascontiguousarray(array)).to(torch.bfloat16) +def raw_volume_with_support(path, target_shape, target_spacing, posterior_shift_mm): + """Same training image, plus original nonzero support before normalization. + + Z-score changes the value of raw air, whereas crop padding stays zero. + Consequently neither the normalized minimum nor normalized zero correctly + identifies both air and padding. Geometry of this mask matches the image. + """ + previous = previous_mr_data_module() + shift = int(round(posterior_shift_mm / target_spacing[2])) + raw = previous.load_and_resample_nii(path, target_spacing) + support = previous.crop_or_pad((raw != 0).astype(np.float32), target_shape, shift) + image = previous.crop_or_pad(previous.NORMALIZERS["zscore"]().normalize(raw), target_shape, shift) + return (torch.from_numpy(np.ascontiguousarray(image)).bfloat16(), + torch.from_numpy(np.ascontiguousarray(support > 0))) + + +def is_mri_sequence(path): + return previous_mr_data_module().is_mri_image_file(Path(path).name) + + class MRAlignedDINO3DDataset(Dataset): """Study-level SSL samples from voxel-aligned, variable-count MR sequences.""" @@ -406,7 +426,7 @@ def __init__( 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 + # MR-RATE 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) @@ -545,7 +565,7 @@ 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 + # MR-RATE 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 diff --git a/dino-pretraining/mr_dino/geometry.py b/dino-pretraining/mr_dino/geometry.py new file mode 100644 index 0000000..5ff94e8 --- /dev/null +++ b/dino-pretraining/mr_dino/geometry.py @@ -0,0 +1,29 @@ +"""MRI sequence kinds and the global-grid Fourier position features.""" +import math +from pathlib import Path +import torch + +SEQUENCE_KINDS=("t1","t1ce","t2","flair","swi","t2star","dwi","adc","pd","stir","tof","other") + + +def sequence_kind(path): + name=Path(path).name.lower() + for token,kind in (("flair","flair"),("swi","swi"),("adc","adc"),("dwi","dwi"), + ("diff","dwi"),("stir","stir"),("tof","tof"),("t2star","t2star"),("t2*","t2star")): + if token in name:return SEQUENCE_KINDS.index(kind) + if "t1" in name: + return SEQUENCE_KINDS.index("t1ce" if any(s in name for s in ("-ce","_ce","post","gad","contrast")) else "t1") + if "t2" in name:return SEQUENCE_KINDS.index("t2") + if "pd" in name:return SEQUENCE_KINDS.index("pd") + return SEQUENCE_KINDS.index("other") + + +def normalized_positions(start,grid,global_grid,device): + axes=[torch.arange(n,device=device,dtype=torch.float32)+s+.5 for s,n in zip(start,grid)] + coordinates=torch.stack(torch.meshgrid(*axes,indexing="ij"),dim=-1).reshape(-1,3) + return coordinates/torch.tensor(global_grid,device=device,dtype=torch.float32) + + +def position_features(coordinates,freqs): + angle=coordinates[:,:,None]*(math.pi*2.**torch.arange(freqs,device=coordinates.device))[None,None,:] + return torch.cat((angle.sin(),angle.cos()),dim=2).reshape(len(coordinates),6*freqs) diff --git a/dino-pretraining/mr_dino/labels.py b/dino-pretraining/mr_dino/labels.py new file mode 100644 index 0000000..151eef1 --- /dev/null +++ b/dino-pretraining/mr_dino/labels.py @@ -0,0 +1,21 @@ +"""Load the final ordered 74-diagnosis MRI label schema.""" +import json +from pathlib import Path +from .probe import read_labels + +DEFAULT_SCHEMA = (Path(__file__).resolve().parents[2] / + "data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/" + "07_mri_diagnosis_extraction/data/mri_diagnoses_74.json") + + +def read_final_labels(labels_csv, schema_path=DEFAULT_SCHEMA): + schema=json.loads(Path(schema_path).read_text()) + names=[item["key"] for item in schema["diagnoses"]] + if len(names)!=74 or len(set(names))!=74: + raise ValueError("Production MRI MIL requires 74 unique diagnosis keys") + columns,labels=read_labels(Path(labels_csv)) + missing=set(names)-set(columns);extra=set(columns)-set(names) + if missing or extra or len(columns)!=74: + raise ValueError(f"Final 74-class label mismatch: missing={sorted(missing)}, extra={sorted(extra)}") + order=[columns.index(name) for name in names] + return names,{uid:values[order] for uid,values in labels.items()} diff --git a/dino-pretraining/mr_dino/mil.py b/dino-pretraining/mr_dino/mil.py new file mode 100644 index 0000000..f97486f --- /dev/null +++ b/dino-pretraining/mr_dino/mil.py @@ -0,0 +1,235 @@ +"""Cache all MRI sequences and train MR-RATE-style final-block + scalar MIL. + +An inference release is initialization, never an optimizer resume. Route masks +are optional external MRI anatomy masks; no CT organ definitions are imported. +""" +from __future__ import annotations + +import argparse +import csv +import json +import math +import os +from pathlib import Path + +import numpy as np +import torch +import torch.distributed as dist + +from .data import raw_volume_with_support, discover_raw_aligned, _load_split_ids, is_mri_sequence +from .mil_head import SparseClassifyThenAggregate +from .mil_stream import LastBlockAdapter, backward, predict +from .probe import auroc, read_labels +from .labels import DEFAULT_SCHEMA, read_final_labels +from .release import load_release, release_metadata + + +def cache(args): + from .preview import encode_volume + rank=int(os.environ.get("RANK",0));world=int(os.environ.get("WORLD_SIZE",1)) + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK",0))) + model,meta=load_release(args.release,"cuda") + cfg=meta["args"];out=Path(args.cache_dir);out.mkdir(parents=True,exist_ok=True) + if args.input_manifest: + rows=json.loads(Path(args.input_manifest).read_text()) + else: + selected=_load_split_ids(args.splits_csv,args.split) + studies=discover_raw_aligned(args.data_folder,selected,cfg["space"]) + rows=[{"study_uid":s["study_uid"],"nifti_files":s["image_paths"]} for s in studies] + tile=tuple(meta["config"]["training_crops"]["global"]) + for row in rows[rank::world]: + uid=row["study_uid"] + if Path(uid).name!=uid:raise ValueError("Unsafe study UID") + complete=out/f"{uid}.study.json" + if complete.exists(): + previous=json.loads(complete.read_text()) + if previous["release_config"]!=meta["config"]:raise ValueError("Cache belongs to another release") + continue + sequences=[];token_count=0 + sources=[f for f in row["nifti_files"] if is_mri_sequence(f)] + if not sources:raise ValueError(f"No MRI sequences for {uid}") + for number,source in enumerate(sources): + volume,support=raw_volume_with_support(source,tuple(cfg["target_shape"]),tuple(cfg["target_spacing"]),cfg["posterior_shift_mm"]) + stem=out/f"{uid}_s{number:03d}" + _,foreground=encode_volume(model,volume,tile,stem,support=support,cache_only=True) + token_count+=int(foreground.sum()) + sequences.append({"cache":stem.with_suffix(".cache.json").name,"source":source}) + pending=complete.with_suffix(".pending") + pending.write_text(json.dumps({"study_uid":uid,"release_config":meta["config"],"sequences":sequences,"token_count":token_count},indent=2)) + pending.replace(complete) + print(uid,"cached",len(sequences),"sequences",flush=True) + + +def study_factory(path,adapter,head,device,return_geometry=False): + study=json.loads(path.read_text());sequences=[] + for sequence in study["sequences"]: + index_path=path.parent/sequence["cache"] + index=json.loads(index_path.read_text()) + if index["dim"]!=adapter.dim or index["prefix"]!=adapter.prefix: + raise ValueError("MRI token cache/model architecture mismatch") + total=sum(c["length"] for c in index["chunks"]) + latent=index_path.parent/index["latent_file"] + if latent.stat().st_size!=total*adapter.dim*2:raise ValueError("Incomplete latent cache") + tokens=np.memmap(latent,mode="r",dtype=np.float16,shape=(total,adapter.dim)) + foreground=np.load(index_path.parent/index["foreground_file"],mmap_mode="r") if "foreground_file" in index else None + sequences.append((index,tokens,foreground,sequence["source"])) + def factory(): + for index,tokens,foreground,source in sequences: + for chunk in index["chunks"]: + start=chunk["offset"];length=chunk["length"] + x=torch.from_numpy(np.array(tokens[start:start+length],copy=True)).to(device,dtype=torch.bfloat16 if device.type=="cuda" else torch.float32) + if foreground is None: + fg=torch.tensor(chunk["foreground"],device=device,dtype=torch.bool) + else: + region=tuple(slice(a,a+n) for a,n in zip(chunk["start"],chunk["grid"])) + fg=torch.tensor(np.asarray(foreground[region]).reshape(-1),device=device,dtype=torch.bool) + with torch.autocast(device.type,dtype=torch.bfloat16,enabled=device.type=="cuda"): + y=adapter(x[None],chunk["grid"])[0,fg] + if not len(y):continue + route=None + if "route_mask" in chunk and head is not None: + head.route_masks_seen=True + route=torch.tensor(chunk["route_mask"],device=device,dtype=torch.bool)[fg] + if route.shape!=(len(y),head.n_classes):raise ValueError("Wrong MRI route dimensions") + if return_geometry: + from .geometry import normalized_positions,sequence_kind + coordinates=normalized_positions(chunk["start"],chunk["grid"],index["token_grid"],device)[fg] + yield y.float(),route,coordinates,sequence_kind(source) + else:yield y.float(),route + return factory + + +def train(args): + local=int(os.environ.get("LOCAL_RANK",0));world=int(os.environ.get("WORLD_SIZE",1)) + torch.cuda.set_device(local);device=torch.device("cuda",local) + if world>1:dist.init_process_group("nccl",device_id=device) + rank=dist.get_rank() if world>1 else 0 + torch.manual_seed(args.seed) + model,meta=load_release(args.release,device if args.mode=="online" else "cpu") + adapter=LastBlockAdapter(model).to(device) + if args.mode=="offline":del model + names,labels=read_final_labels(args.labels_csv,args.label_schema) + normalizer=torch.load(args.normalization,map_location=device,weights_only=False) + if normalizer['release_config']!=meta['config'] or normalizer['split']!='train':raise ValueError('Wrong normalization provenance') + mean=normalizer['mean'].float();rms=float(normalizer['rms']) + if mean.shape!=(adapter.dim,) or not math.isfinite(rms) or rms<=0:raise ValueError('Invalid MRI normalization') + head=SparseClassifyThenAggregate(adapter.dim+6*args.position_freqs,len(names),use_norm=False,checkpoint_chunks=False).to(device) + # Soft MRI routes, when provided. Missing routes mean all tokens are eligible. + head.route_bias_raw=torch.nn.Parameter(torch.full((len(names),),2.,device=device)) + params=list(adapter.parameters())+list(head.parameters()) + optimizer=torch.optim.AdamW([{"params":adapter.parameters(),"lr":args.block_lr}, + {"params":[p for n,p in head.named_parameters() if n!="route_bias_raw"],"lr":args.head_lr}, + {"params":[head.route_bias_raw],"lr":args.head_lr,"weight_decay":0.}],weight_decay=args.weight_decay) + with open(args.splits_csv,newline="") as handle:rows=list(csv.DictReader(handle)) + splits={r["study_uid"]:r["split"] for r in rows} + patient={r["study_uid"]:r["patient_uid"] for r in rows} + patient_sets={s:{patient[u] for u in splits if splits[u]==s} for s in ("train","val","test")} + if any(patient_sets[a]&patient_sets[b] for a,b in (("train","val"),("train","test"),("val","test"))): + raise ValueError("Patient leakage between MR-RATE splits") + caches={} + if args.mode=="offline": + if not args.cache_dir:raise ValueError("Offline mode requires --cache-dir") + for path in sorted(Path(args.cache_dir).glob("*.study.json")): + row=json.loads(path.read_text());uid=row["study_uid"] + if row["release_config"]!=meta["config"]:raise ValueError("Mixed checkpoints in token cache") + if uid in caches:raise ValueError("Duplicate study cache") + caches[uid]=path + else: + if not args.data_folder:raise ValueError("Online mode requires --data-folder") + for study in discover_raw_aligned(args.data_folder,set(splits),meta["args"]["space"]): + if study["study_uid"] in caches:raise ValueError("Duplicate online study") + caches[study["study_uid"]]=study + def make_factory(uid): + if args.mode=="offline":raw=study_factory(caches[uid],adapter,head,device,return_geometry=True) + else: + from .vision import online_study_factory + raw=online_study_factory(caches[uid],model,adapter,meta,device,return_geometry=True) + def transformed(): + from .geometry import position_features + for y,mask,coordinates,_ in raw(): + yield torch.cat(((y-mean)/rms,position_features(coordinates,args.position_freqs)*args.position_scale),1),mask + return transformed + train_ids=sorted(u for u in caches if splits.get(u)=="train" and u in labels) + val_ids=sorted(u for u in caches if splits.get(u)=="val" and u in labels) + if not train_ids or not val_ids:raise ValueError("Both train and validation caches are required") + missing=[u for u in labels if splits.get(u) in ('train','val') and u not in caches] + if missing and not args.allow_incomplete_data:raise ValueError(f'Incomplete MRI cohort: {len(missing)} labelled train/val studies missing; first={missing[:10]}') + out=Path(args.out);out.mkdir(parents=True,exist_ok=True) + start=0 + if args.resume: + state=torch.load(args.resume,map_location=device,weights_only=False) + if state["release_config"]!=meta["config"] or state["classes"]!=names:raise ValueError("Resume checkpoint mismatch") + adapter.load_state_dict(state["adapter"]);head.load_state_dict(state["head"]) + optimizer.load_state_dict(state["optimizer"]);start=state["epoch"] + def objective(logits,target): + known=torch.isfinite(target)&(target>=0) + if not bool(known.any()):raise ValueError("Study has no known labels") + if not bool(((target[known]==0)|(target[known]==1)).all()):raise ValueError("Nonbinary pathology labels") + return torch.nn.functional.binary_cross_entropy_with_logits(logits[known],target[known]) + for epoch in range(start,args.epochs): + order=torch.randperm(len(train_ids),generator=torch.Generator().manual_seed(args.seed+epoch)).tolist() + steps=math.ceil(len(order)/world) + order=order+[-1]*(steps*world-len(order)) + head.train();adapter.train();losses=[] + for index in order[rank::world]: + optimizer.zero_grad(set_to_none=True) + real=index>=0 + if real: + uid=train_ids[index] + target=torch.tensor(labels[uid],device=device) + loss,_=backward(head,make_factory(uid),target,objective) + losses.append(float(loss)) + if world>1: + count=torch.tensor(float(real),device=device) + dist.all_reduce(count) + for param in params: + if param.grad is None:param.grad=torch.zeros_like(param) + dist.all_reduce(param.grad);param.grad.div_(count.clamp_min(1)) + torch.nn.utils.clip_grad_norm_(params,3.,error_if_nonfinite=True) + optimizer.step() + head.eval();adapter.eval();predictions=[] + for uid in val_ids[rank::world]: + probability=predict(head,make_factory(uid)).sigmoid()[0].cpu().tolist() + predictions.append({"study_uid":uid,"probabilities":probability}) + if world>1: + gathered=[None]*world;dist.all_gather_object(gathered,predictions) + predictions=[r for shard in gathered for r in shard] + meter=torch.tensor([sum(losses),len(losses)],device=device,dtype=torch.float64) + if world>1:dist.all_reduce(meter) + routed=torch.tensor(int(getattr(head,"route_masks_seen",False)),device=device) + if world>1:dist.all_reduce(routed,op=dist.ReduceOp.MAX) + if rank==0: + per_class={} + for c,name in enumerate(names): + valid=[r for r in predictions if np.isfinite(labels[r["study_uid"]][c]) and labels[r["study_uid"]][c]>=0] + per_class[name]=auroc(np.array([labels[r["study_uid"]][c] for r in valid]),np.array([r["probabilities"][c] for r in valid])) if valid else None + values=[v for v in per_class.values() if v is not None] + result={"epoch":epoch+1,"train_loss":float(meter[0]/meter[1].clamp_min(1)),"macro_auroc":float(np.mean(values)) if values else None,"per_class_auroc":per_class,"predictions":predictions} + (out/f"epoch_{epoch+1:03d}.json").write_text(json.dumps(result,indent=2)) + state={"format":"mrdino_mil_v1","epoch":epoch+1,"release_config":meta["config"],"classes":names,"adapter":adapter.state_dict(),"head":head.state_dict(),"optimizer":optimizer.state_dict(),"arguments":vars(args),"normalization":normalizer,"architecture":{"dim":head.dim,"n_classes":len(names),"hidden_dim":512,"mlp_hidden_dim":384,"use_norm":False},"route_masks_used":bool(routed)} + pending=out/"checkpoint.pending";torch.save(state,pending);pending.replace(out/f"epoch_{epoch+1:03d}.pt") + print({k:v for k,v in result.items() if k not in {"predictions","per_class_auroc"}},flush=True) + if world>1:dist.barrier(device_ids=[local]) + if world>1:dist.destroy_process_group() + + +def main(): + parser=argparse.ArgumentParser();sub=parser.add_subparsers(dest="command",required=True) + c=sub.add_parser("cache");c.add_argument("--release",required=True);c.add_argument("--cache-dir",required=True) + source=c.add_mutually_exclusive_group(required=True);source.add_argument("--input-manifest");source.add_argument("--data-folder") + c.add_argument("--splits-csv");c.add_argument("--split",default="train",choices=["train","val","test"]) + t=sub.add_parser("train");t.add_argument("--release",required=True);t.add_argument("--cache-dir") + t.add_argument("--mode",default="offline",choices=["offline","online"]);t.add_argument("--data-folder") + t.add_argument('--normalization',required=True);t.add_argument('--position-freqs',type=int,default=4);t.add_argument('--position-scale',type=float,default=1.) + t.add_argument('--allow-incomplete-data',action='store_true',help='Explicit implementation smoke only; production requires complete labelled train/val coverage') + t.add_argument("--labels-csv",required=True);t.add_argument("--splits-csv",required=True);t.add_argument("--out",required=True) + t.add_argument('--label-schema',default=str(DEFAULT_SCHEMA),help='Final ordered 74-diagnosis MRI schema') + t.add_argument("--head-lr",type=float,required=True);t.add_argument("--block-lr",type=float,required=True) + t.add_argument("--weight-decay",type=float,default=.01);t.add_argument("--epochs",type=int,default=10) + t.add_argument("--seed",type=int,default=42);t.add_argument("--resume") + args=parser.parse_args() + if args.command=="cache":cache(args) + else:train(args) + + +if __name__=="__main__":main() diff --git a/dino-pretraining/mr_dino/mil_head.py b/dino-pretraining/mr_dino/mil_head.py new file mode 100644 index 0000000..502ffbb --- /dev/null +++ b/dino-pretraining/mr_dino/mil_head.py @@ -0,0 +1,526 @@ +"""Class-specific gated-attention MIL for multi-sequence MRI patch tokens. + +A study contains every eligible sequence. Tokens are classified individually, +then pooled with a separate attention softmax for each diagnosis. Ragged bag +boundaries keep studies independent. Sparse aggregation streams token chunks. +""" + +from __future__ import annotations + +from typing import Sequence, Tuple, Union + +import torch +from torch import nn +from torch.nn import functional as F + + +__all__ = [ + "PatchClassifier", + "ClassifyThenAggregate", + "SparseClassifyThenAggregate", + "MRRATEMILHead", + "collate_ragged_bags", +] + + +def _positive_int(value: int, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + return value + + +class PatchClassifier(nn.Module): + """Shared instance classifier: Linear/GELU blocks followed by Linear.""" + + def __init__( + self, + dim: int, + hidden_dims: Sequence[int], + out_dim: int, + ) -> None: + super().__init__() + dim = _positive_int(dim, "dim") + out_dim = _positive_int(out_dim, "out_dim") + hidden_dims = tuple( + _positive_int(hidden, f"hidden_dims[{index}]") + for index, hidden in enumerate(hidden_dims) + ) + + layers = [] + current_dim = dim + for hidden_dim in hidden_dims: + layers.extend((nn.Linear(current_dim, hidden_dim), nn.GELU())) + current_dim = hidden_dim + layers.append(nn.Linear(current_dim, out_dim)) + self.net = nn.Sequential(*layers) + + def forward(self, instances: torch.Tensor) -> torch.Tensor: + return self.net(instances) + + +class ClassifyThenAggregate(nn.Module): + """class-specific class-specific gated-attention MIL head. + + Parameters + ---------- + dim: + Width of each visual-instance embedding. + n_classes: + Number of independent study-level labels. + hidden_dim: + Width of each gated-attention projection. + mlp_hidden_dims: + Hidden widths in the shared instance classifier. An empty sequence + makes the instance classifier a single linear layer. + drop_rate: + Dropout applied independently before the attention and classifier + branches. + use_gating: + If true, attention is ``tanh(Vx) * sigmoid(Ux)``. If false it is + ``tanh(Vx)``. + use_norm: + Apply a separate LayerNorm before each branch. + use_output_bias_scale: + Learn a per-class affine transform of the aggregated logits. + + Notes + ----- + ``forward(..., return_details=True)`` returns + ``(bag_logits, class_attention, patch_logits)``. Both detail tensors are + flat ``[N, C]`` tensors aligned with the input instance order. For every bag + and class, ``class_attention`` sums to one. + """ + + def __init__( + self, + dim: int, + n_classes: int, + hidden_dim: int = 512, + mlp_hidden_dims: Sequence[int] = (384,), + drop_rate: float = 0.0, + init_std: float = 0.02, + use_gating: bool = True, + use_norm: bool = False, + use_output_bias_scale: bool = True, + ) -> None: + super().__init__() + self.dim = _positive_int(dim, "dim") + self.n_classes = _positive_int(n_classes, "n_classes") + hidden_dim = _positive_int(hidden_dim, "hidden_dim") + if not 0.0 <= drop_rate < 1.0: + raise ValueError(f"drop_rate must be in [0, 1), got {drop_rate}") + if init_std <= 0.0: + raise ValueError(f"init_std must be positive, got {init_std}") + + self.use_gating = bool(use_gating) + self.use_output_bias_scale = bool(use_output_bias_scale) + self.init_std = float(init_std) + + def make_norm() -> nn.Module: + if use_norm: + return nn.LayerNorm(self.dim, eps=1e-6) + return nn.Identity() + + self.attention_norm = make_norm() + self.classifier_norm = make_norm() + self.input_dropout = nn.Dropout(drop_rate) + + self.attention_v = nn.Linear(self.dim, hidden_dim) + self.attention_u = nn.Linear(self.dim, hidden_dim) if self.use_gating else None + self.attention_w = nn.Linear(hidden_dim, self.n_classes) + self.patch_classifier = PatchClassifier( + self.dim, + mlp_hidden_dims, + self.n_classes, + ) + + if self.use_output_bias_scale: + self.output_scale = nn.Parameter(torch.ones(self.n_classes)) + self.output_bias = nn.Parameter(torch.zeros(self.n_classes)) + else: + self.register_parameter("output_scale", None) + self.register_parameter("output_bias", None) + + self.apply(self._init_weights) + + @property + def input_dim(self) -> int: + """Alias useful when inspecting generic visual-instance heads.""" + + return self.dim + + @property + def num_classes(self) -> int: + """Alias useful when inspecting generic visual-instance heads.""" + + return self.n_classes + + def _init_weights(self, module: nn.Module) -> None: + if isinstance(module, nn.Linear): + nn.init.trunc_normal_(module.weight, std=self.init_std) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.LayerNorm): + nn.init.ones_(module.weight) + nn.init.zeros_(module.bias) + + def _validate_inputs( + self, + instances: torch.Tensor, + bag_ptr: torch.Tensor, + ) -> Tuple[int, ...]: + if not isinstance(instances, torch.Tensor): + raise TypeError("instances must be a torch.Tensor") + if instances.ndim != 2 or instances.shape[1] != self.dim: + raise ValueError( + f"instances must have shape [N, {self.dim}], " + f"got {tuple(instances.shape)}" + ) + if not instances.is_floating_point(): + raise TypeError("instances must have a floating-point dtype") + + if not isinstance(bag_ptr, torch.Tensor): + raise TypeError("bag_ptr must be a torch.Tensor") + if bag_ptr.ndim != 1: + raise ValueError( + f"bag_ptr must be one-dimensional, got {tuple(bag_ptr.shape)}" + ) + if bag_ptr.dtype not in ( + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + ): + raise TypeError("bag_ptr must have an integer dtype") + + bounds = tuple(int(bound) for bound in bag_ptr.detach().cpu().tolist()) + if len(bounds) < 2: + raise ValueError("bag_ptr must contain at least a start and end boundary") + if bounds[0] != 0: + raise ValueError("bag_ptr must start at 0") + if bounds[-1] != instances.shape[0]: + raise ValueError( + "bag_ptr must end at the number of flat instances " + f"({instances.shape[0]}), got {bounds[-1]}" + ) + if any(end <= start for start, end in zip(bounds[:-1], bounds[1:])): + raise ValueError( + "bag_ptr boundaries must be strictly increasing; " + "MIL bags must contain at least one instance" + ) + return bounds + + def forward( + self, + instances: torch.Tensor, + bag_ptr: torch.Tensor, + return_details: bool = False, + ) -> Union[ + torch.Tensor, + Tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ]: + """Classify visual instances and aggregate each ragged study bag. + + Parameters + ---------- + instances: + Concatenated visual-instance embeddings with shape ``[N, dim]``. + bag_ptr: + Integer cumulative boundaries with shape ``[B + 1]``. It must + begin at zero, end at ``N``, and be strictly increasing. + return_details: + Also return flat class attention and instance-level logits. + """ + + bounds = self._validate_inputs(instances, bag_ptr) + + attention_input = self.input_dropout(self.attention_norm(instances)) + classifier_input = self.input_dropout(self.classifier_norm(instances)) + + attention_hidden = torch.tanh(self.attention_v(attention_input)) + if self.attention_u is not None: + attention_hidden = attention_hidden * torch.sigmoid( + self.attention_u(attention_input) + ) + attention_scores = self.attention_w(attention_hidden) + patch_logits = self.patch_classifier(classifier_input) + + bag_logits = [] + attention_chunks = [] + for start, end in zip(bounds[:-1], bounds[1:]): + class_attention = F.softmax(attention_scores[start:end], dim=0) + bag_logits.append((class_attention * patch_logits[start:end]).sum(dim=0)) + if return_details: + attention_chunks.append(class_attention) + + output = torch.stack(bag_logits, dim=0) + if self.output_scale is not None: + output = output * self.output_scale + self.output_bias + + if return_details: + return output, torch.cat(attention_chunks, dim=0), patch_logits + return output + + +class SparseClassifyThenAggregate(nn.Module): + """Class-selective token MIL with exact chunked softmax aggregation. + + Unmentioned pathology leaves are unknown, not negative. ``class_indices`` + therefore selects only the explicitly supervised classes for the current + report while parameters still span the complete vocabulary. An online + stable-softmax reduction avoids materializing ``all_tokens x classes``. + """ + + def __init__( + self, + dim: int, + n_classes: int, + hidden_dim: int = 512, + mlp_hidden_dim: int = 384, + drop_rate: float = 0.0, + use_norm: bool = False, + token_chunk_size: int = 32768, + checkpoint_chunks: bool = True, + ) -> None: + super().__init__() + self.dim = _positive_int(dim, "dim") + self.n_classes = _positive_int(n_classes, "n_classes") + self.hidden_dim = _positive_int(hidden_dim, "hidden_dim") + self.mlp_hidden_dim = _positive_int(mlp_hidden_dim, "mlp_hidden_dim") + self.token_chunk_size = _positive_int(token_chunk_size, "token_chunk_size") + if not 0.0 <= drop_rate < 1.0: + raise ValueError("drop_rate must be in [0, 1)") + self.checkpoint_chunks = bool(checkpoint_chunks) + self.attention_norm = ( + nn.LayerNorm(self.dim, eps=1e-6) if use_norm else nn.Identity() + ) + self.classifier_norm = ( + nn.LayerNorm(self.dim, eps=1e-6) if use_norm else nn.Identity() + ) + self.input_dropout = nn.Dropout(drop_rate) + self.attention_v = nn.Linear(self.dim, self.hidden_dim) + self.attention_u = nn.Linear(self.dim, self.hidden_dim) + self.attention_w = nn.Linear(self.hidden_dim, self.n_classes) + self.patch_hidden = nn.Linear(self.dim, self.mlp_hidden_dim) + self.patch_out = nn.Linear(self.mlp_hidden_dim, self.n_classes) + self.output_scale = nn.Parameter(torch.ones(self.n_classes)) + self.output_bias = nn.Parameter(torch.zeros(self.n_classes)) + self.apply(self._init_weights) + + def _init_weights(self, module: nn.Module) -> None: + if isinstance(module, nn.Linear): + nn.init.trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.LayerNorm): + nn.init.ones_(module.weight) + nn.init.zeros_(module.bias) + + def _chunk_outputs( + self, tokens: torch.Tensor, class_indices: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + attention_scores = self._chunk_attention_scores(tokens, class_indices) + classifier_input = self.input_dropout(self.classifier_norm(tokens)) + patch_hidden = F.gelu(self.patch_hidden(classifier_input)) + patch_logits = F.linear( + patch_hidden, + self.patch_out.weight.index_select(0, class_indices), + self.patch_out.bias.index_select(0, class_indices), + ) + return attention_scores, patch_logits + + def _chunk_attention_scores( + self, tokens: torch.Tensor, class_indices: torch.Tensor + ) -> torch.Tensor: + """Class-attention scores without the unused patch-logit branch.""" + attention_input = self.input_dropout(self.attention_norm(tokens)) + attention_hidden = torch.tanh(self.attention_v(attention_input)) + attention_hidden = attention_hidden * torch.sigmoid( + self.attention_u(attention_input) + ) + attention_scores = F.linear( + attention_hidden, + self.attention_w.weight.index_select(0, class_indices), + self.attention_w.bias.index_select(0, class_indices), + ) + return attention_scores + + def forward( + self, tokens: torch.Tensor, class_indices: torch.Tensor + ) -> torch.Tensor: + if tokens.ndim != 2 or tokens.shape[1] != self.dim or tokens.shape[0] == 0: + raise ValueError( + f"tokens must be nonempty [N,{self.dim}], got {tuple(tokens.shape)}" + ) + if class_indices.ndim != 1 or class_indices.numel() == 0: + raise ValueError("class_indices must be a nonempty vector") + class_indices = class_indices.to(device=tokens.device, dtype=torch.long) + if ( + int(class_indices.min()) < 0 + or int(class_indices.max()) >= self.n_classes + or class_indices.unique().numel() != class_indices.numel() + ): + raise ValueError("class_indices are out of range or duplicated") + + from torch.utils.checkpoint import checkpoint + + running_max = torch.full( + (class_indices.numel(),), + -torch.inf, + dtype=tokens.dtype, + device=tokens.device, + ) + running_denominator = torch.zeros_like(running_max) + running_numerator = torch.zeros_like(running_max) + for start in range(0, tokens.shape[0], self.token_chunk_size): + chunk = tokens[start : start + self.token_chunk_size] + if self.checkpoint_chunks and self.training and torch.is_grad_enabled(): + scores, patch_logits = checkpoint( + self._chunk_outputs, + chunk, + class_indices, + use_reentrant=False, + ) + else: + scores, patch_logits = self._chunk_outputs(chunk, class_indices) + chunk_max = scores.max(dim=0).values + new_max = torch.maximum(running_max, chunk_max) + old_scale = torch.exp(running_max - new_max) + weights = torch.exp(scores - new_max.unsqueeze(0)) + running_numerator = ( + running_numerator * old_scale + + (weights * patch_logits).sum(dim=0) + ) + running_denominator = ( + running_denominator * old_scale + weights.sum(dim=0) + ) + running_max = new_max + logits = running_numerator / running_denominator.clamp_min( + torch.finfo(running_denominator.dtype).tiny + ) + logits = ( + logits * self.output_scale.index_select(0, class_indices) + + self.output_bias.index_select(0, class_indices) + ) + return logits.unsqueeze(0) + + def forward_with_evidence( + self, tokens: torch.Tensor, class_indices: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return bag logits and exact class-attention-pooled token features. + + The evidence vector for class ``c`` is + + ``sum_t softmax(attention_score[t, c]) * tokens[t]``. + + It is accumulated in the same online-softmax pass as the class logit, + so every token contributes without materializing an ``N x C x D`` + tensor. This method is intended for frozen-MIL conditioning; the + ordinary :meth:`forward` path remains unchanged for MIL training. + """ + if tokens.ndim != 2 or tokens.shape[1] != self.dim or tokens.shape[0] == 0: + raise ValueError( + f"tokens must be nonempty [N,{self.dim}], got {tuple(tokens.shape)}" + ) + if class_indices.ndim != 1 or class_indices.numel() == 0: + raise ValueError("class_indices must be a nonempty vector") + class_indices = class_indices.to(device=tokens.device, dtype=torch.long) + if ( + int(class_indices.min()) < 0 + or int(class_indices.max()) >= self.n_classes + or class_indices.unique().numel() != class_indices.numel() + ): + raise ValueError("class_indices are out of range or duplicated") + + running_max = torch.full( + (class_indices.numel(),), + -torch.inf, + dtype=tokens.dtype, + device=tokens.device, + ) + running_denominator = torch.zeros_like(running_max) + running_logit_numerator = torch.zeros_like(running_max) + # Evidence accumulation is small (C x D) and benefits from fp32 when + # a study contains hundreds of thousands of uncapped image tokens. + running_evidence_numerator = torch.zeros( + (class_indices.numel(), self.dim), + dtype=torch.float32, + device=tokens.device, + ) + for start in range(0, tokens.shape[0], self.token_chunk_size): + chunk = tokens[start : start + self.token_chunk_size] + scores, patch_logits = self._chunk_outputs(chunk, class_indices) + chunk_max = scores.max(dim=0).values + new_max = torch.maximum(running_max, chunk_max) + old_scale = torch.exp(running_max - new_max) + weights = torch.exp(scores - new_max.unsqueeze(0)) + running_logit_numerator = ( + running_logit_numerator * old_scale + + (weights * patch_logits).sum(dim=0) + ) + running_evidence_numerator = ( + running_evidence_numerator * old_scale.float().unsqueeze(1) + + weights.float().transpose(0, 1) @ chunk.float() + ) + running_denominator = ( + running_denominator * old_scale + weights.sum(dim=0) + ) + running_max = new_max + + denominator = running_denominator.clamp_min( + torch.finfo(running_denominator.dtype).tiny + ) + logits = running_logit_numerator / denominator + logits = ( + logits * self.output_scale.index_select(0, class_indices) + + self.output_bias.index_select(0, class_indices) + ) + evidence = ( + running_evidence_numerator / denominator.float().unsqueeze(1) + ).to(tokens.dtype) + return logits.unsqueeze(0), evidence.unsqueeze(0) + + +# MRI head alias. +MRRATEMILHead = ClassifyThenAggregate + + +def collate_ragged_bags( + batch: Sequence[Tuple[torch.Tensor, torch.Tensor]], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Collate ``(visual_instances, target)`` samples without padding. + + Returns ``(flat_features, targets, bag_ptr)``. Boundaries stay on CPU, + matching PyTorch's usual collation metadata; the head accepts boundaries on + either CPU or the feature device. + """ + + if len(batch) == 0: + raise ValueError("cannot collate an empty batch") + + bags, targets = zip(*batch) + first = bags[0] + if not isinstance(first, torch.Tensor) or first.ndim != 2: + raise ValueError("each visual-instance bag must be a rank-2 tensor [M, D]") + feature_dim = first.shape[1] + + for index, bag in enumerate(bags): + if not isinstance(bag, torch.Tensor) or bag.ndim != 2: + raise ValueError(f"bag {index} must be a rank-2 tensor [M, D]") + if bag.shape[0] == 0: + raise ValueError(f"bag {index} is empty") + if bag.shape[1] != feature_dim: + raise ValueError( + f"bag {index} has feature width {bag.shape[1]}, " + f"expected {feature_dim}" + ) + if not bag.is_floating_point(): + raise TypeError(f"bag {index} must have a floating-point dtype") + + if not all(isinstance(target, torch.Tensor) for target in targets): + raise TypeError("each target must be a torch.Tensor") + + lengths = torch.tensor([bag.shape[0] for bag in bags], dtype=torch.long) + bag_ptr = torch.cat((torch.zeros(1, dtype=torch.long), lengths.cumsum(0))) + return torch.cat(bags, dim=0), torch.stack(targets, dim=0), bag_ptr diff --git a/dino-pretraining/mr_dino/mil_stream.py b/dino-pretraining/mr_dino/mil_stream.py new file mode 100644 index 0000000..ca73e86 --- /dev/null +++ b/dino-pretraining/mr_dino/mil_stream.py @@ -0,0 +1,108 @@ +"""MR-RATE scalar MIL, with the MRI checkpoint's own final block and physical RoPE. + +Every sequence/crop contributes to one study-level class softmax. Two replayed +passes give exact gradients without retaining every crop's transformer graph. +""" +from __future__ import annotations + +import copy +import torch +from torch import nn + + +class LastBlockAdapter(nn.Module): + def __init__(self, backbone): + super().__init__() + self.block = copy.deepcopy(backbone.blocks[-1]).float() + self.norm = copy.deepcopy(backbone.norm).float() + self.block.requires_grad_(True) + self.norm.requires_grad_(True) + self.rope = copy.deepcopy(backbone.rope_embed).float() + self.prefix = 1 + backbone.n_storage_tokens + self.dim = backbone.embed_dim + self.rope.requires_grad_(False) + + def forward(self, tokens, grid): + # Disable pretraining DropPath for deterministic replay. Parameters train. + self.block.eval() + rope = tuple(t.to(device=tokens.device, dtype=tokens.dtype) + for t in self.rope(D=grid[0], H=grid[1], W=grid[2])) + return self.norm(self.block(tokens, rope))[:, self.prefix:] + + +def _scores(head, scores, mask): + scores = scores.float() + if mask is not None: + if mask.dtype != torch.bool or mask.shape != scores.shape: + raise ValueError("Route mask must be boolean [tokens, classes]") + penalty = torch.nn.functional.softplus(head.route_bias_raw.float()) + scores = scores - (~mask).float() * penalty + return scores + + +@torch.no_grad() +def statistics(head, factory): + device = next(head.parameters()).device + classes = torch.arange(head.n_classes, device=device) + maximum = torch.full((head.n_classes,), -torch.inf, device=device) + den = torch.zeros_like(maximum) + num = torch.zeros_like(maximum) + count = 0 + for tokens, mask in factory(): + scores, logits = head._chunk_outputs(tokens.float(), classes) + scores = _scores(head, scores, mask) + new = torch.maximum(maximum, scores.amax(0)) + old = (maximum-new).exp() + weights = (scores-new).exp() + den = den*old + weights.sum(0) + num = num*old + (weights*logits.float()).sum(0) + maximum = new + count += len(tokens) + if count == 0 or not bool((den > 0).all()) or not bool(torch.isfinite(num).all()): + raise ValueError("Empty or nonfinite MIL bag") + return maximum, den, num/den, count + + +@torch.no_grad() +def predict(head, factory): + _, _, pooled, _ = statistics(head, factory) + return (pooled*head.output_scale + head.output_bias)[None] + + +def backward(head, factory, target, loss_fn=None): + """Replay deterministic features; gradients also flow into the final block.""" + if head.input_dropout.p: + raise ValueError("Two-pass replay requires dropout=0") + versions = tuple(p._version for p in head.parameters()) + maximum, den, pooled, count = statistics(head, factory) + leaf = pooled.detach().requires_grad_(True) + logits = (leaf*head.output_scale+head.output_bias)[None] + objective = loss_fn or torch.nn.functional.binary_cross_entropy_with_logits + loss = objective(logits, target.reshape(1,-1)) + loss.backward() + gradient = leaf.grad.detach() + classes = torch.arange(head.n_classes, device=pooled.device) + seen = 0 + for tokens, mask in factory(): + scores, patch = head._chunk_outputs(tokens.float(), classes) + scores = _scores(head, scores, mask) + with torch.no_grad(): + attention = (scores.detach()-maximum).exp()/den + dl = attention*gradient + ds = dl*(patch.float()-pooled) + torch.autograd.backward((scores,patch), (ds,dl)) + seen += len(tokens) + if seen != count or versions != tuple(p._version for p in head.parameters()): + raise RuntimeError("Token count or head parameters changed during replay") + return loss.detach(), logits.detach() + + +@torch.no_grad() +def evidence(head, factory): + """Yield exact positive token logit contributions; bias is excluded.""" + maximum, den, _, _ = statistics(head, factory) + classes = torch.arange(head.n_classes, device=maximum.device) + for tokens, mask in factory(): + scores, patch = head._chunk_outputs(tokens.float(),classes) + attention = (_scores(head,scores,mask)-maximum).exp()/den + yield (attention*patch.float()*head.output_scale).clamp_min(0) diff --git a/dino-pretraining/mr_dino/model.py b/dino-pretraining/mr_dino/model.py index 4d4e972..216febb 100644 --- a/dino-pretraining/mr_dino/model.py +++ b/dino-pretraining/mr_dino/model.py @@ -21,7 +21,7 @@ except ImportError as exc: # pragma: no cover - actionable cluster error raise ImportError( "Official DINOv3 is required. Set DINOV3_ROOT and add it to PYTHONPATH; " - "see scripts/ctdino3d/README.md." + "see dino-pretraining/README.md." ) from exc diff --git a/dino-pretraining/mr_dino/objective.py b/dino-pretraining/mr_dino/objective.py index 9680e4a..774f943 100644 --- a/dino-pretraining/mr_dino/objective.py +++ b/dino-pretraining/mr_dino/objective.py @@ -45,7 +45,7 @@ def __init__( nlayers=3, ) if normalize_prototypes: - # FORA fix: cosine prototypes let the teacher commit from step 0 (plain std-0.02 + # MR-RATE 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) @@ -264,7 +264,7 @@ def reset_gram_anchor_from_teacher(self) -> None: 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 + """Cross-crop patch objective (MR-RATE 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 @@ -347,7 +347,7 @@ def weighted_mean(per_sample: Tensor) -> Tensor: 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. + # MR-RATE 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), @@ -444,7 +444,7 @@ 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 + # Teacher-target diagnostics (MR-RATE 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) diff --git a/dino-pretraining/mr_dino/preview.py b/dino-pretraining/mr_dino/preview.py new file mode 100644 index 0000000..c00b4bd --- /dev/null +++ b/dino-pretraining/mr_dino/preview.py @@ -0,0 +1,134 @@ +"""Whole-volume three-component PCA previews, with exact tiled token geometry.""" +from __future__ import annotations + +import argparse +import html +import json +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import torch +import torch.nn.functional as F +from sklearn.decomposition import PCA + +from .data import raw_volume_with_support, is_mri_sequence +from .release import load_release + + +@torch.inference_mode() +def encode_volume(model, volume, tile, output, patch=(2, 16, 16), support=None, cache_only=False): + if support is None or tuple(support.shape) != tuple(volume.shape): + raise ValueError("Original MRI support mask is required") + background = float(volume.float().min()) + shape = tuple(volume.shape) + pads = [] + for n, t in reversed(list(zip(shape, tile))): + pads.extend((0, (-n) % t)) + padded = F.pad(volume.float(), pads, value=0) + valid = F.pad(support.float(), pads, value=0) + grid = tuple(n // p for n, p in zip(padded.shape, patch)) + features = None if cache_only else np.zeros((*grid, model.embed_dim), dtype=np.float16) + foreground = np.zeros(grid, dtype=bool) + chunks = [] + latent_path = output.with_suffix(".latents.f16") + tile_grid = tuple(t // p for t, p in zip(tile, patch)) + offset = 0 + with latent_path.open("wb") as sink: + for z in range(0, padded.shape[0], tile[0]): + for y in range(0, padded.shape[1], tile[1]): + for x in range(0, padded.shape[2], tile[2]): + crop = padded[z:z+tile[0], y:y+tile[1], x:x+tile[2]][None, None].cuda().bfloat16() + supported = valid[z:z+tile[0], y:y+tile[1], x:x+tile[2]][None,None].cuda() + fg = F.max_pool3d(supported, patch, patch)[0,0] > 0 + if not bool(fg.any()): + continue + tokens, g = model.prepare_tokens_with_masks(crop) + rope = tuple(c.to(device=tokens.device, dtype=tokens.dtype) + for c in model.rope_embed(D=g[0], H=g[1], W=g[2])) + with torch.autocast("cuda", dtype=torch.bfloat16): + for block in model.blocks[:-1]: + tokens = block(tokens, rope) + latent = tokens.to(dtype=torch.float16).cpu().numpy()[0] + sink.write(latent.tobytes()) + if not cache_only: + tokens = model.blocks[-1](tokens, rope) + patch_tokens = model.norm(tokens)[:, 1+model.n_storage_tokens:] + start = tuple(v//p for v, p in zip((z,y,x), patch)) + region = tuple(slice(a,a+n) for a,n in zip(start,tile_grid)) + if not cache_only:features[region] = patch_tokens.reshape(*tile_grid, model.embed_dim).float().cpu().numpy() + foreground[region] = fg.cpu().numpy() + chunks.append({"offset": offset, "length": len(latent), "grid": list(g), + "start": list(start)}) + offset += len(latent) + output.with_suffix(".cache.json").write_text(json.dumps({ + "format": "mrdino_penultimate_tiles_v1", "dim": model.embed_dim, + "prefix": 1+model.n_storage_tokens, "latent_file": latent_path.name, + "foreground_file": output.with_suffix(".foreground.npy").name, + "foreground_tokens": int(foreground.sum()), + "support_rule": "nonzero original resampled MRI before zscore; original training image unchanged", + "chunks": chunks, "token_grid": grid, "volume_shape": shape})) + if not cache_only:np.save(output.with_suffix(".tokens.npy"), features) + np.save(output.with_suffix(".foreground.npy"), foreground) + return features, foreground + + +def render(volume, features, foreground, out, title): + values = features[foreground].astype(np.float32) + if len(values) < 3: + raise ValueError("Fewer than three foreground tokens") + pca = PCA(n_components=3, svd_solver="randomized", random_state=0) + scores = pca.fit_transform(values) + lo, hi = np.percentile(scores, [1,99], axis=0) + rgb = np.zeros((*foreground.shape,3), np.float32) + rgb[foreground] = np.clip((scores-lo)/(hi-lo+1e-8),0,1) + np.savez(out.with_suffix(".pca.npz"), mean=pca.mean_, components=pca.components_, + explained_variance_ratio=pca.explained_variance_ratio_, rgb=rgb, + lower=lo, upper=hi) + image = volume.float().numpy() + spacing=np.array([1.,.5,.5]);patch=np.array([2,16,16]);pitch=spacing*patch + center = np.median(np.argwhere(foreground),axis=0).astype(int) + fig, axes = plt.subplots(2,3,figsize=(13,8),facecolor="#111827") + for axis in range(3): + sl = [slice(None)]*3; sl[axis] = int(center[axis]); sl=tuple(sl) + image_sl=[slice(None)]*3 + image_sl[axis]=min(int(center[axis]*patch[axis]+patch[axis]//2),image.shape[axis]-1) + remaining=[i for i in range(3) if i!=axis] + axes[0,axis].imshow(image[tuple(image_sl)].T,cmap="gray",origin="lower",aspect=spacing[remaining[1]]/spacing[remaining[0]]) + axes[1,axis].imshow(rgb[sl].transpose(1,0,2),origin="lower",aspect=pitch[remaining[1]]/pitch[remaining[0]]) + axes[0,axis].set_title(f"Input: axis {axis}, index {center[axis]}",color="white") + axes[1,axis].set_title("Whole-volume PCA → RGB",color="white") + for row in range(2): axes[row,axis].axis("off") + fig.suptitle(title,color="white");fig.tight_layout();fig.savefig(out.with_suffix(".png"),dpi=140);plt.close(fig) + + +def main(): + p=argparse.ArgumentParser() + p.add_argument("--release",required=True);p.add_argument("--manifest",required=True) + p.add_argument("--out",required=True) + args=p.parse_args();out=Path(args.out);out.mkdir(parents=True,exist_ok=True) + torch.cuda.set_device(0) + model,meta=load_release(args.release,"cuda") + train=meta["args"];tile=tuple(meta["config"]["training_crops"]["global"]) + panels=[] + for row in json.loads(Path(args.manifest).read_text()): + files=[f for f in row["nifti_files"] if is_mri_sequence(f)] + uid=row["study_uid"];sequences=[] + for number,source in enumerate(files): + stem=f"{uid}_s{number:03d}" + print(uid,"preprocessing",source,flush=True) + volume,support=raw_volume_with_support(source,tuple(train["target_shape"]),tuple(train["target_spacing"]),train["posterior_shift_mm"]) + features,fg=encode_volume(model,volume,tile,out/stem,support=support) + render(volume,features,fg,out/stem,uid+" · "+Path(source).name) + sequences.append({"cache":stem+".cache.json","source":source}) + panels.append(f'

{html.escape(uid)}

{html.escape(Path(source).name)}

') + print(uid,"PCA saved",int(fg.sum()),"foreground tokens",flush=True) + (out/"index.html").write_text('MR DINO embeddings

MR DINO: whole-volume PCA

Three-component PCA fitted independently to all foreground tokens of each volume. Non-overlapping tiles assembled before PCA. Colors are independent between sequences. Prototype heads are absent from the inference releases.

'+''.join(panels)) + (out/f"{uid}.study.json").write_text(json.dumps({"study_uid":uid,"release_config":meta["config"],"sequences":sequences},indent=2)) + (out/"provenance.json").write_text(json.dumps(meta,indent=2)) + (out/"index.html").write_text('MR DINO embeddings

MR DINO: whole-volume PCA

Three-component PCA fitted independently to all foreground tokens of each volume. Non-overlapping inference tiles are assembled before PCA. Colors are independent between scans. Learned prototype heads are absent from these inference releases.

'+''.join(panels)) + + +if __name__=="__main__":main() diff --git a/dino-pretraining/mr_dino/probe.py b/dino-pretraining/mr_dino/probe.py index 9108072..b0f0377 100644 --- a/dino-pretraining/mr_dino/probe.py +++ b/dino-pretraining/mr_dino/probe.py @@ -1,4 +1,4 @@ -"""Validation MIL probe for MR DINO checkpoints (the FORA CT-DINO checkpoint protocol). +"""Validation MIL probe for MR DINO checkpoints (the MR DINO checkpoint protocol). A checkpoint is scored without touching train or test data: @@ -14,7 +14,7 @@ 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 +3. ``cv``: patient-level K-fold cross-validation of the 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. @@ -57,6 +57,9 @@ 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() and (checkpoint / "config.json").is_file(): + from .release import release_metadata + return release_metadata(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)} @@ -71,6 +74,9 @@ def load_teacher_backbone(checkpoint: Path, device: torch.device | str = "cpu"): from .train_ddp import build_backbone checkpoint = Path(checkpoint) + if checkpoint.is_dir() and (checkpoint / "config.json").is_file(): + from .release import load_release + return load_release(checkpoint, device) meta = read_checkpoint_metadata(checkpoint) backbone = build_backbone(meta["args"]["arch"]) if checkpoint.is_dir(): @@ -268,7 +274,7 @@ def extract(args) -> Path: # --------------------------------------------------------------------------- MIL head def mil_head_class(): - """NeuroVFM ClassifyThenAggregate from MR-RATE's MIL probe (single source of truth).""" + """MRI 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) diff --git a/dino-pretraining/mr_dino/recipe.py b/dino-pretraining/mr_dino/recipe.py index 08c65d2..6d1b9c9 100644 --- a/dino-pretraining/mr_dino/recipe.py +++ b/dino-pretraining/mr_dino/recipe.py @@ -1,4 +1,4 @@ -"""Recipe fixes ported from the FORA CT DINOv3 pipeline (v13/v14, Sep 2026). +"""Recipe fixes ported from the MR 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: diff --git a/dino-pretraining/mr_dino/release.py b/dino-pretraining/mr_dino/release.py new file mode 100644 index 0000000..d800a94 --- /dev/null +++ b/dino-pretraining/mr_dino/release.py @@ -0,0 +1,67 @@ +"""Load Dong's inference releases without confusing them with DCP resumes.""" +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path + +import torch + + +def release_metadata(path: str | Path) -> dict: + path = Path(path).resolve() + config = json.loads((path / "config.json").read_text()) + if config.get("source") != "model.teacher.backbone": + raise ValueError("Expected an EMA teacher backbone release") + provenance = "release training metadata" + checkpoint_metadata = path / "training/checkpoint_metadata.json" + if checkpoint_metadata.is_file(): + args = json.loads(checkpoint_metadata.read_text())["args"] + else: + candidates = [path / "training/run_config.json", + path / f"training/run_config_{config['stage']}.json"] + found = next((p for p in candidates if p.is_file()), None) + if found is not None: + args = json.loads(found.read_text()) + elif config.get("format")=="mrdino3d_7b_teacher_backbone_bf16_v1" and config.get("mr_rate_commit")=="43c05c6": + # This legacy release omitted the run arguments. Its README states + # atlas space; the pinned trainer below supplies these defaults. + # NVIDIA can supply the original overrides in preprocessing.json. + config.setdefault("space","atlas_space") + config.setdefault("training_crops",{"global":config["backbone"]["volume_size"]}) + args={"arch":config["arch"],"space":"atlas_space","target_shape":[256,384,384], + "target_spacing":[1.,.5,.5],"posterior_shift_mm":15.} + provenance="legacy 7B README + pinned 43c05c61bf98d00c4e825a27a0fca3c742754d78 trainer defaults; original run overrides absent" + if (path/"preprocessing.json").is_file(): + args.update(json.loads((path/"preprocessing.json").read_text())) + provenance="legacy release + supplied preprocessing.json" + else: + raise ValueError("Release has no training preprocessing metadata; supply its training configuration") + for key in ("arch", "space"): + if args.get(key) != config[key]: + raise ValueError(f"Release/training metadata disagree on {key}") + if tuple(args["target_spacing"]) != tuple(config["backbone"]["voxel_spacing_mm"]): + raise ValueError("Release/training voxel spacing mismatch") + weights = (path / config["weights"]).resolve() + if not weights.is_relative_to(path): + raise ValueError("Weights must belong to the release directory") + return {"step": config["step"], "stage": config["stage"], "args": args, + "source": str(path), "config": config,"preprocessing_provenance":provenance} + + +def load_release(path: str | Path, device="cpu"): + """Use the supplied model source, strict keys, and memory-mapped BF16 weights.""" + path = Path(path).resolve() + metadata = release_metadata(path) + config = metadata["config"] + spec = importlib.util.spec_from_file_location("dong_release_model", path / "mr_dino/model.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + constructor = getattr(module, f"vit_{config['arch']}_3d") + with torch.device("meta"): + backbone = constructor(**config["backbone"]) + state = torch.load(path / config["weights"], map_location="cpu", weights_only=True, mmap=True) + backbone.load_state_dict(state, strict=True, assign=True) + backbone.to(device=device) + backbone.eval().requires_grad_(False) + return backbone, metadata diff --git a/dino-pretraining/mr_dino/train_ddp.py b/dino-pretraining/mr_dino/train_ddp.py index 7f1accb..75d8960 100644 --- a/dino-pretraining/mr_dino/train_ddp.py +++ b/dino-pretraining/mr_dino/train_ddp.py @@ -111,7 +111,7 @@ def parse_args() -> argparse.Namespace: 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) + # MR-RATE 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)) diff --git a/dino-pretraining/mr_dino/train_fsdp.py b/dino-pretraining/mr_dino/train_fsdp.py index aab63f0..ea28cb1 100644 --- a/dino-pretraining/mr_dino/train_fsdp.py +++ b/dino-pretraining/mr_dino/train_fsdp.py @@ -71,7 +71,7 @@ def parse_args() -> argparse.Namespace: 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", "7b"), default="hplus", - help="hplus = ViT-H+ 0.84B (MR default); large = ViT-L 0.30B (FORA CT); 7b kept for reference") + help="hplus = ViT-H+ 0.84B (MR default); large = ViT-L 0.30B (MRI); 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( @@ -124,7 +124,7 @@ def parse_args() -> argparse.Namespace: p.add_argument("--compile", 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) + # MR-RATE 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") diff --git a/dino-pretraining/mr_dino/vision.py b/dino-pretraining/mr_dino/vision.py new file mode 100644 index 0000000..9e56fad --- /dev/null +++ b/dino-pretraining/mr_dino/vision.py @@ -0,0 +1,45 @@ +"""Image-backed, replayable MRI features using the same geometry as cached MIL.""" +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from .data import raw_volume_with_support + + +def online_study_factory(study, backbone, adapter, metadata, device,return_geometry=False): + args=metadata["args"] + tile=tuple(metadata["config"]["training_crops"]["global"]) + patch=tuple(metadata["config"]["backbone"]["patch_size"]) + volumes=[raw_volume_with_support(source,tuple(args["target_shape"]), + tuple(args["target_spacing"]),args["posterior_shift_mm"]) + for source in study["image_paths"]] + backbone.eval().requires_grad_(False) + def factory(): + for sequence,(volume,support) in enumerate(volumes): + pads=[] + for n,t in reversed(list(zip(volume.shape,tile))):pads.extend((0,(-n)%t)) + image=F.pad(volume,pads,value=0) + valid=F.pad(support.float(),pads,value=0) + for z in range(0,image.shape[0],tile[0]): + for y in range(0,image.shape[1],tile[1]): + for x in range(0,image.shape[2],tile[2]): + sl=(slice(z,z+tile[0]),slice(y,y+tile[1]),slice(x,x+tile[2])) + keep=F.max_pool3d(valid[sl][None,None],patch,patch).flatten()>0 + if not bool(keep.any()):continue + with torch.no_grad(),torch.autocast(device.type,dtype=torch.bfloat16): + crop=image[sl][None,None].to(device) + tokens,grid=backbone.prepare_tokens_with_masks(crop) + rope=tuple(r.to(device=tokens.device,dtype=tokens.dtype) + for r in backbone.rope_embed(D=grid[0],H=grid[1],W=grid[2])) + for block in backbone.blocks[:-1]:tokens=block(tokens,rope) + with torch.autocast(device.type,dtype=torch.bfloat16): + features=adapter(tokens,grid)[0,keep.to(device)] + if return_geometry: + from .geometry import normalized_positions,sequence_kind + start=tuple(a//p for a,p in zip((z,y,x),patch)) + global_grid=tuple(n//p for n,p in zip(image.shape,patch)) + coordinates=normalized_positions(start,grid,global_grid,device)[keep.to(device)] + yield features.float(),None,coordinates,sequence_kind(study["image_paths"][sequence]) + else:yield features.float(),None + return factory diff --git a/dino-pretraining/mr_dino/vlm.py b/dino-pretraining/mr_dino/vlm.py new file mode 100644 index 0000000..8092ac1 --- /dev/null +++ b/dino-pretraining/mr_dino/vlm.py @@ -0,0 +1,256 @@ +"""the two-pass VL-CABS training, ported to multi-sequence MR DINO. + +Frozen prefix caches are shared with MRI MIL. The final block starts from DINO, +not from a MIL checkpoint. Online mode recomputes the same frozen prefix. +""" +from __future__ import annotations +import argparse, contextlib, csv, hashlib, json, math, os, random, signal +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace +import numpy as np +import torch +import torch.distributed as dist +from .data import discover_raw_aligned +from .geometry import position_features +from .mil import study_factory +from .mil_stream import LastBlockAdapter +from .release import load_release +from .vision import online_study_factory +from .vlm_core import TextEncoder,VisualProjector,streamed_step + + +def setup(): + world=int(os.environ.get('WORLD_SIZE',1));rank=int(os.environ.get('RANK',0));local=int(os.environ.get('LOCAL_RANK',0)) + cpu=os.environ.get('MRRATE_VLM_DEVICE')=='cpu' + device=torch.device('cpu') if cpu else torch.device('cuda',local) + if not cpu:torch.cuda.set_device(local) + if world>1:dist.init_process_group('gloo' if cpu else 'nccl',device_id=None if cpu else device,timeout=timedelta(minutes=30)) + return rank,world,local,device + + +def barrier(device): + if dist.is_initialized(): + dist.barrier() if device.type=='cpu' else dist.barrier(device_ids=[device.index]) + + +def read_sentences(path): + result={} + with open(path) as stream: + for line_number,line in enumerate(stream,1): + if line.startswith('version https://git-lfs'): + raise ValueError('Sentences are a Git LFS pointer; run git lfs pull for this file') + row=json.loads(line) + if not row.get('valid_json',False):continue + sentences=row.get('extracted_sentences',[]) + if not sentences:continue + uid=str(row['volume_name']) + if uid in result:raise ValueError(f'Duplicate sentence study at line {line_number}: {uid}') + if not all(isinstance(s,str) and s.strip() for s in sentences):raise ValueError(f'Invalid sentences for {uid}') + result[uid]=sentences + if not result:raise ValueError('No usable finding sentences') + return result + + +def backend(config,device): + online=config.get('mode','offline')=='online' + backbone,metadata=load_release(config['release'],device if online else 'cpu') + adapter=LastBlockAdapter(backbone).to(device) + with open(config['splits_csv'],newline='') as stream:rows=list(csv.DictReader(stream)) + membership={r['study_uid']:r['split'] for r in rows} + patients={s:{r['patient_uid'] for r in rows if r['split']==s} for s in ('train','val','test')} + if any(patients[a]&patients[b] for a,b in (('train','val'),('train','test'),('val','test'))):raise ValueError('Patient split leakage') + if online: + studies={s['study_uid']:s for s in discover_raw_aligned(config['data_folder'],set(membership),metadata['args']['space'])} + else: + del backbone + studies={} + for path in sorted(Path(config['cache_dir']).glob('*.study.json')): + row=json.loads(path.read_text());uid=row['study_uid'] + if row['release_config']!=metadata['config']:raise ValueError('Cache belongs to another DINO checkpoint') + if uid in studies:raise ValueError('Duplicate MRI cache') + studies[uid]=path + def factory(uid): + if online:return online_study_factory(studies[uid],backbone,adapter,metadata,device,return_geometry=True) + return study_factory(studies[uid],adapter,None,device,return_geometry=True) + if not studies:raise ValueError('No MR studies available') + if dist.is_initialized(): + signature=hashlib.sha256(json.dumps({'ids':sorted(studies),'release':metadata['config']},sort_keys=True).encode()).hexdigest() + signatures=[None]*dist.get_world_size();dist.all_gather_object(signatures,signature) + if len(set(signatures))!=1:raise ValueError('Ranks see different MRI cohorts/releases; use a shared complete cache or image tree') + return adapter,metadata,membership,studies,factory + + +def normalize(config,output,rank,world,device): + adapter,metadata,membership,studies,factory=backend(config,device) + ids=sorted(u for u in studies if membership.get(u)=='train') + missing=[u for u in membership if membership[u]=='train' and u not in studies] + if missing and not config.get('allow_incomplete_data',False):raise ValueError(f'Incomplete training normalization cohort: {len(missing)} missing MRI studies') + sums=torch.zeros(adapter.dim,device=device,dtype=torch.float64) + squared=torch.zeros((),device=device,dtype=torch.float64);count=torch.zeros_like(squared) + with torch.no_grad(): + for uid in ids[rank::world]: + for y,_,_,_ in factory(uid)(): + sums+=y.double().sum(0);squared+=y.double().square().sum();count+=len(y) + if world>1: + for value in (sums,squared,count):dist.all_reduce(value) + if count<=0:raise ValueError('Empty training normalization population') + mean=sums/count + rms=((squared/count-mean.square().sum())/adapter.dim).clamp_min(1e-12).sqrt() + if rank==0: + path=Path(output);path.parent.mkdir(parents=True,exist_ok=True) + state={'mean':mean.float().cpu(),'rms':float(rms),'release_config':metadata['config'], + 'split':'train','studies':len(ids),'tokens':int(count)} + pending=path.with_suffix('.pending');torch.save(state,pending);pending.replace(path) + print({'event':'normalization_complete','studies':len(ids),'tokens':int(count),'rms':float(rms)},flush=True) + barrier(device) + + +def batches(ids,world,seed=None): + order=list(ids) + if seed is not None:random.Random(seed).shuffle(order) + for start in range(0,len(order),world):yield order[start:start+world]+[None]*max(0,start+world-len(order)) + + +def train(config,output,rank,world,local,device,max_steps=0): + if world<2:raise ValueError('MR-RATE VLM uses one study per rank and needs at least two distributed ranks for negatives') + cfg=config['training'];spec=cfg['model'];seed=cfg['seed'];torch.manual_seed(seed) + adapter,metadata,membership,studies,raw_factory=backend(config,device) + sentences=read_sentences(config['sentences']) + train_ids=sorted(u for u in studies if membership.get(u)=='train' and u in sentences) + val_ids=sorted(u for u in studies if membership.get(u)=='val' and u in sentences) + if not train_ids or not val_ids:raise ValueError('Training and validation MRI/sentence cohorts are required') + missing=[u for u in sentences if membership.get(u) in ('train','val') and u not in studies] + if missing and not config.get('allow_incomplete_data',False):raise ValueError(f'Incomplete VLM cohort: {len(missing)} train/val MRI studies missing') + norm=torch.load(config['normalization'],map_location=device,weights_only=False) + if norm['release_config']!=metadata['config'] or norm['split']!='train':raise ValueError('Wrong normalization provenance') + mean=norm['mean'].float();rms=float(norm['rms']) + if mean.shape!=(adapter.dim,) or not math.isfinite(rms) or rms<=0:raise ValueError('Invalid MRI normalization') + frequencies=cfg.get('position_freqs',4);dim=spec['dim'] + visual=VisualProjector(adapter.dim+6*frequencies,dim).to(device) + text=TextEncoder(spec['text_model'],dim).to(device) + from transformers import BertTokenizer + tokenizer=BertTokenizer.from_pretrained(spec['text_model'],do_lower_case=True) + temperature=torch.nn.Parameter(torch.tensor(math.log(1/.07),device=device)) + pool_temperature=torch.nn.Parameter(torch.tensor(math.log(spec.get('pool_temperature_init',1/.07)),device=device)) + groups=[{'params':[p for p in adapter.parameters() if p.requires_grad and p.ndim>1],'lr':spec['lr_block'],'weight_decay':spec['wd']}, + {'params':[p for p in adapter.parameters() if p.requires_grad and p.ndim<=1],'lr':spec['lr_block'],'weight_decay':0.}, + {'params':visual.parameters(),'lr':spec['lr_proj'],'weight_decay':spec['wd']}, + {'params':text.bert.parameters(),'lr':spec['lr_text'],'weight_decay':spec['wd']}, + {'params':text.proj.parameters(),'lr':spec['lr_proj'],'weight_decay':spec['wd']}, + {'params':[temperature,pool_temperature],'lr':spec['lr_temp'],'weight_decay':0.}] + optimizer=torch.optim.AdamW(groups) + parameters=[p for group in optimizer.param_groups for p in group['params']] + base_lrs=[g['lr'] for g in optimizer.param_groups] + out=Path(output);out.mkdir(parents=True,exist_ok=True) + epoch0=batch0=step=0;best=float('inf') + latest=out/'latest.pt' + if latest.exists(): + state=torch.load(latest,map_location=device,weights_only=False) + if state['config']!=config or state['world']!=world or state['release_config']!=metadata['config']:raise ValueError('Resume recipe/world/release mismatch') + adapter.load_state_dict(state['adapter']);visual.load_state_dict(state['visual']);text.load_state_dict(state['text']);optimizer.load_state_dict(state['optimizer']) + with torch.no_grad():temperature.copy_(state['temperature']);pool_temperature.copy_(state['pool_temperature']) + epoch0,batch0,step,best=state['epoch'],state['batch'],state['step'],state['best_val_loss'] + rng=state['rng'][rank];torch.set_rng_state(rng['torch'].cpu()) + if device.type=='cuda':torch.cuda.set_rng_state(rng['cuda'].cpu(),device) + random.setstate(rng['python']);np.random.set_state(rng['numpy']) + else: + torch.manual_seed(seed+rank);random.seed(seed+rank);np.random.seed(seed+rank) + stopping=False;requested=False + def stop(*_): + nonlocal stopping + stopping=True + def request(*_): + nonlocal requested + requested=True + signal.signal(signal.SIGTERM,stop);signal.signal(signal.SIGUSR1,request) + def save(epoch,batch,name='latest.pt'): + rng={'torch':torch.get_rng_state(),'cuda':torch.cuda.get_rng_state(device) if device.type=='cuda' else None, + 'python':random.getstate(),'numpy':np.random.get_state()} + all_rng=[None]*world;dist.all_gather_object(all_rng,rng) + if rank==0: + state={'format':'mrdino_vlm_v1','config':config,'world':world,'release_config':metadata['config'], + 'epoch':epoch,'batch':batch,'step':step,'best_val_loss':best,'adapter':adapter.state_dict(), + 'visual':visual.state_dict(),'text':text.state_dict(),'temperature':temperature.detach(), + 'pool_temperature':pool_temperature.detach(),'optimizer':optimizer.state_dict(),'rng':all_rng} + pending=out/(name+'.pending');torch.save(state,pending);pending.replace(out/name) + barrier(device) + def visual_factory(uid): + if uid is None: + def empty():yield visual(torch.zeros(1,adapter.dim+6*frequencies,device=device),0) + return empty + raw=raw_factory(uid) + def chunks(): + for y,_,coordinates,kind in raw(): + features=torch.cat(((y-mean)/rms,position_features(coordinates,frequencies)*spec.get('position_scale',1.)),dim=1) + yield visual(features,kind).float() + return chunks + slots=cfg['max_sentences'];total=math.ceil(len(train_ids)/world)*cfg['epochs'] + def run_case(uid,training): + values=list(sentences[uid]) if uid else [] + chooser=random.Random(seed*7919+step*131+rank) + if len(values)>slots:values=chooser.sample(values,slots) + n=len(values);values+=['']*(slots-n) + mask=torch.arange(slots,device=device)=max_steps)),device=device) + request_flag=torch.tensor(int(requested),device=device) + dist.all_reduce(stop_flag,op=dist.ReduceOp.MAX);dist.all_reduce(request_flag,op=dist.ReduceOp.MAX) + if rank==0: + progress={'epoch':epoch+1,'batch':batch+1,'step':step,'loss':float(loss),'lr_fraction':fraction, + 'gradient_norm':norm_value,'rank0_tokens':count,'temperature':float(temperature.detach().clamp(0,4).exp()), + 'pool_temperature':float(pool_temperature.detach().clamp(0,4).exp())} + (out/'progress.json').write_text(json.dumps(progress,indent=2));print(progress,flush=True) + if step%cfg.get('save_every',200)==0 or stop_flag or request_flag: + save(epoch,batch+1);requested=False + if stop_flag:return + adapter.eval();visual.eval();text.eval();numerator=denominator=correct=0. + with torch.no_grad(): + for group in batches(val_ids,world): + loss,logits,mask,_=run_case(group[rank],False) + valid=int(mask.sum());numerator+=float(loss)*valid;denominator+=valid + expected=torch.arange(world,device=device).repeat_interleave(slots) + correct+=int(((logits.argmax(1)==expected)&mask).sum()) + val_loss=numerator/max(denominator,1) + improved=val_loss per-sentence max m, normaliser S and pooled vector p = sum_k a_k z_k with + a_k = softmax_k(pool_temp * cos(t, z_k)). Pass 2 (grad): for each recomputed chunk the surrogate + J = sum_ik a_ik.detach() * (g_i . z_k) + sum_ik s_ik * [a_ik (g_i . z_k - g_i . p_i)].detach() + back-propagates exactly dL/dz_k = a_k g + (dL/ds_k) ds_k/dz_k and dL/dt, dL/dpool_temp of the attention path, given g = dL/dp.""" + def __init__(self,text,pool_temp): + self.t=text.detach();n=len(self.t);D_=self.t.shape[1];dev=self.t.device + dt=self.t.dtype;self.pt=float(pool_temp);self.m=torch.full((n,),-torch.inf,device=dev,dtype=dt);self.S=torch.zeros(n,device=dev,dtype=dt);self.P=torch.zeros(n,D_,device=dev,dtype=dt) + @torch.no_grad() + def add(self,z): + s=self.pt*(self.t@z.T)*z.norm(dim=-1).clamp_min(1e-12).reciprocal()[None] + mn=torch.maximum(self.m,s.amax(1));c=torch.exp(self.m-mn);e=torch.exp(s-mn[:,None]) + self.S=self.S*c+e.sum(1);self.P=self.P*c[:,None]+e@z;self.m=mn + def pooled(self):return self.P/self.S[:,None] + def surrogate(self,t_leaf,pool_temp,z,p,g): + s=pool_temp*(t_leaf@z.T)*z.norm(dim=-1).clamp_min(1e-12).reciprocal()[None] + with torch.no_grad(): + a=torch.exp(s-self.m[:,None])/self.S[:,None];gz=g@z.T;ds=a*(gz-(g*p).sum(1,keepdim=True)) + return (a*(g@z.T)).sum()+(s*ds).sum() + +def pooled_logits_reference(text,z,pool_temp,temp): + """Plain autograd version (whole z in memory): the definition the streaming path must reproduce (tests only).""" + w=torch.softmax(pool_temp*((text@z.T)*z.norm(dim=-1).clamp_min(1e-12).reciprocal()[None]),dim=1);pooled=w@z + return (F.normalize(pooled,dim=-1)*text).sum(1)*temp + +class GatherLayer(torch.autograd.Function): + """All-gather with a DDP-safe, deterministic backward. + + torch.distributed.nn.functional.all_gather issues its backward collective + (reduce_scatter) through the autograd-distributed machinery, whose op + ordering interleaves nondeterministically with DDP's gradient all_reduces + on the shared communicator -> NCCL same-communicator order mismatch that + deadlocks backward at scale (observed: 64 ranks hang right after step 0). + + This is the battle-tested OpenCLIP pattern: forward does a plain blocking + all_gather; backward does a single blocking all_reduce(SUM) of the stacked + grads and returns this rank's slice. The collective sits at one fixed + autograd node with identical ordering on every rank. + """ + @staticmethod + def forward(ctx, x): + x = x.contiguous() + out = [torch.empty_like(x) for _ in range(dist.get_world_size())] + dist.all_gather(out, x) + return tuple(out) + + @staticmethod + def backward(ctx, *grads): + grad_stack = torch.stack(grads, dim=0) + dist.all_reduce(grad_stack, op=dist.ReduceOp.SUM) + return grad_stack[dist.get_rank()] + +def contrastive_ce(logits, num_sentences_per_image, mask_all, device): + """Symmetric (text<->image) InfoNCE used by both the global and local + terms. Each sentence's positive is its own study; negatives are other + studies (and, for image->text, other studies' sentences).""" + logits = logits.float() + B_global = logits.shape[1] + N_T = logits.shape[0] + sentence_to_image_map = torch.arange(B_global, device=device).repeat_interleave(num_sentences_per_image) + image_membership = F.one_hot(sentence_to_image_map, num_classes=B_global).bool() + + if mask_all is not None: + mask_float = mask_all.float() + loss_t_all = F.cross_entropy(logits, sentence_to_image_map, reduction='none') + loss_t_all = torch.where(mask_all, loss_t_all, torch.zeros_like(loss_t_all)) + loss_t = loss_t_all.sum() / mask_float.sum().clamp(min=1.0) + + pos_logits = logits[torch.arange(N_T, device=device), sentence_to_image_map] + not_from_image = ~image_membership + valid_sentences = mask_all.view(N_T, 1).expand(N_T, B_global) + logits_masked_for_neg = logits.clone() + logits_masked_for_neg[~(not_from_image & valid_sentences)] = -1e4 + logsumexp_neg = torch.logsumexp(logits_masked_for_neg, dim=0)[sentence_to_image_map] + log_denom = torch.logsumexp(torch.stack([pos_logits, logsumexp_neg], dim=1), dim=1) + loss_i_all = -pos_logits + log_denom + loss_i_all = torch.where(mask_all, loss_i_all, torch.zeros_like(loss_i_all)) + loss_i = loss_i_all.sum() / mask_float.sum().clamp(min=1.0) + else: + loss_t = F.cross_entropy(logits, sentence_to_image_map) + pos_logits = logits[torch.arange(N_T, device=device), sentence_to_image_map] + logits_masked_for_neg = logits.clone() + logits_masked_for_neg[~(~image_membership)] = -1e4 + logsumexp_neg = torch.logsumexp(logits_masked_for_neg, dim=0)[sentence_to_image_map] + log_denom = torch.logsumexp(torch.stack([pos_logits, logsumexp_neg], dim=1), dim=1) + loss_i = (-pos_logits + log_denom).mean() + return loss_t + loss_i diff --git a/dino-pretraining/requirements.txt b/dino-pretraining/requirements.txt index dbc6686..891b697 100644 --- a/dino-pretraining/requirements.txt +++ b/dino-pretraining/requirements.txt @@ -1,3 +1,7 @@ numpy nibabel torch>=2.6 +transformers>=4.51 +matplotlib +scikit-learn +pytest diff --git a/dino-pretraining/scripts/smoke_mil_vlm.sh b/dino-pretraining/scripts/smoke_mil_vlm.sh new file mode 100755 index 0000000..6efabaf --- /dev/null +++ b/dino-pretraining/scripts/smoke_mil_vlm.sh @@ -0,0 +1,24 @@ +#!/bin/bash +set -euo pipefail +export OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MRRATE_VLM_DEVICE=cpu TOKENIZERS_PARALLELISM=false +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CODE="$(dirname "$SCRIPT_DIR")" +: "${DINOV3_ROOT:?Set DINOV3_ROOT to the pinned official DINOv3 checkout}" +export PYTHONPATH="$CODE:$DINOV3_ROOT:${PYTHONPATH:-}" +R=${1:-$(mktemp -d /tmp/mrrate-vlm-smoke.XXXXXX)} +python3 "$CODE/tests/make_vlm_lifecycle_fixture.py" "$R" +python3 -m torch.distributed.run --standalone --nproc-per-node=2 -m mr_dino.vlm normalize --config "$R/config.json" --output "$R/normalization.pt" +python3 -m torch.distributed.run --standalone --nproc-per-node=2 -m mr_dino.vlm train --config "$R/config.json" --output "$R/uninterrupted" +python3 -m torch.distributed.run --standalone --nproc-per-node=2 -m mr_dino.vlm train --config "$R/config.json" --output "$R/resumed" --max-steps 1 +python3 -m torch.distributed.run --standalone --nproc-per-node=2 -m mr_dino.vlm train --config "$R/config.json" --output "$R/resumed" +python3 - "$R" <<'PY' +import sys,torch,json +from pathlib import Path +r=Path(sys.argv[1]);a=torch.load(r/'uninterrupted/latest.pt',weights_only=False);b=torch.load(r/'resumed/latest.pt',weights_only=False) +for group in ['adapter','visual','text']: + for key in a[group]:torch.testing.assert_close(a[group][key],b[group][key],rtol=0,atol=0) +for key in ['temperature','pool_temperature']:torch.testing.assert_close(a[key],b[key],rtol=0,atol=0) +assert a['step']==b['step']==3 and a['epoch']==b['epoch']==1 +result={'synthetic_implementation_test':True,'world':2,'training_studies':5,'validation_studies':2,'steps':3,'includes_padded_tail':True,'normalization':'train only','checkpoint_resume':'bitwise identical final adapter, projector, BERT and temperatures'} +(r/'result.json').write_text(json.dumps(result,indent=2));print(json.dumps(result)) +PY diff --git a/dino-pretraining/scripts/train_32n_7b.sbatch b/dino-pretraining/scripts/train_32n_7b.sbatch index 8a41253..69c8e2e 100755 --- a/dino-pretraining/scripts/train_32n_7b.sbatch +++ b/dino-pretraining/scripts/train_32n_7b.sbatch @@ -1,5 +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. +# LEGACY recipe (DINOv3-7B continuation settings). The MRI 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 3b2c1af..d470ba9 100755 --- a/dino-pretraining/scripts/train_32n_hplus.sbatch +++ b/dino-pretraining/scripts/train_32n_hplus.sbatch @@ -1,5 +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. +# LEGACY recipe (DINOv3-7B continuation settings). The MRI 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 index edd3f2f..71027b5 100644 --- a/dino-pretraining/scripts/train_mrdino.sbatch +++ b/dino-pretraining/scripts/train_mrdino.sbatch @@ -41,7 +41,7 @@ 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 +# Phases 2/3 restart the schedule at step 0, so they need their own settings (MR 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. @@ -50,7 +50,7 @@ fi # 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). +# Defaults keep the MR-RATE 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) @@ -81,7 +81,7 @@ 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" +echo "MR DINOv3 $ARCH (MR-RATE 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" \ diff --git a/dino-pretraining/tests/make_vlm_lifecycle_fixture.py b/dino-pretraining/tests/make_vlm_lifecycle_fixture.py new file mode 100644 index 0000000..8a27551 --- /dev/null +++ b/dino-pretraining/tests/make_vlm_lifecycle_fixture.py @@ -0,0 +1,32 @@ +import json,sys +from pathlib import Path +import numpy as np,torch +from transformers import BertConfig,BertModel +from mr_dino.model import DinoVisionTransformer3D +root=Path(sys.argv[1]);root.mkdir(parents=True,exist_ok=True) +release=root/'release';(release/'training').mkdir(parents=True,exist_ok=True);(release/'mr_dino').mkdir(exist_ok=True) +backbone_args={'volume_size':[4,32,32],'patch_size':[2,16,16],'voxel_spacing_mm':[1.,.5,.5],'in_chans':1,'n_storage_tokens':2,'drop_path_rate':0.} +config={'format':'test_teacher','arch':'tiny','source':'model.teacher.backbone','space':'atlas_space','step':4,'stage':'highres','weights':'teacher.pt','backbone':backbone_args,'training_crops':{'global':[4,32,32]}} +(release/'config.json').write_text(json.dumps(config)) +args={'arch':'tiny','space':'atlas_space','target_shape':[4,32,32],'target_spacing':[1.,.5,.5],'posterior_shift_mm':0.} +(release/'training/checkpoint_metadata.json').write_text(json.dumps({'args':args})) +(release/'mr_dino/model.py').write_text('from mr_dino.model import DinoVisionTransformer3D\ndef vit_tiny_3d(**kwargs):\n return DinoVisionTransformer3D(embed_dim=32,depth=2,num_heads=4,**kwargs)\n') +torch.manual_seed(12);model=DinoVisionTransformer3D(embed_dim=32,depth=2,num_heads=4,**backbone_args);model.init_weights();model.eval();torch.save(model.state_dict(),release/'teacher.pt') +cache=root/'cache';cache.mkdir(exist_ok=True);rows=[];sentences=[] +for i in range(7): + uid=f's{i}';split='train' if i<5 else 'val';rows.append(f'b,p{i},{uid},{split}') + with torch.no_grad(): + x=torch.randn(1,1,4,32,32);tokens,grid=model.prepare_tokens_with_masks(x) + rope=model.rope_embed(D=grid[0],H=grid[1],W=grid[2]) + tokens=model.blocks[0](tokens,rope)[0].numpy().astype(np.float16) + (cache/(uid+'.f16')).write_bytes(tokens.tobytes());np.save(cache/(uid+'.fg.npy'),np.ones(grid,dtype=bool)) + index={'format':'mrdino_penultimate_tiles_v1','dim':32,'prefix':3,'latent_file':uid+'.f16','foreground_file':uid+'.fg.npy','token_grid':list(grid),'chunks':[{'offset':0,'length':len(tokens),'start':[0,0,0],'grid':list(grid)}]} + (cache/(uid+'.cache.json')).write_text(json.dumps(index)) + (cache/(uid+'.study.json')).write_text(json.dumps({'study_uid':uid,'release_config':config,'sequences':[{'cache':uid+'.cache.json','source':uid+'_t1w.nii.gz'}]})) + sentences.append({'volume_name':uid,'valid_json':True,'extracted_sentences':['alpha beta','gamma delta'] if i%2 else ['beta gamma','delta alpha']}) +(root/'split.csv').write_text('batch_id,patient_uid,study_uid,split\n'+'\n'.join(rows)+'\n');(root/'sentences.jsonl').write_text(''.join(json.dumps(r)+'\n' for r in sentences)) +text=root/'bert';text.mkdir(exist_ok=True) +vocab=['[PAD]','[UNK]','[CLS]','[SEP]','[MASK]','alpha','beta','gamma','delta'];(text/'vocab.txt').write_text('\n'.join(vocab)+'\n');(text/'tokenizer_config.json').write_text(json.dumps({'tokenizer_class':'BertTokenizer','do_lower_case':True})) +BertModel(BertConfig(vocab_size=len(vocab),hidden_size=16,num_hidden_layers=1,num_attention_heads=2,intermediate_size=24,max_position_embeddings=64)).save_pretrained(text) +recipe={'mode':'offline','release':str(release),'cache_dir':str(cache),'splits_csv':str(root/'split.csv'),'sentences':str(root/'sentences.jsonl'),'normalization':str(root/'normalization.pt'),'training':{'epochs':1,'seed':17,'warmup_steps':0,'max_sentences':2,'max_tokens':16,'save_every':1,'position_freqs':2,'model':{'dim':8,'text_model':str(text),'lr_block':1e-4,'lr_proj':1e-3,'lr_text':1e-4,'lr_temp':1e-3,'wd':.01}}} +(root/'config.json').write_text(json.dumps(recipe));print('Lifecycle fixture ready',flush=True) diff --git a/dino-pretraining/tests/test_mil_stream.py b/dino-pretraining/tests/test_mil_stream.py new file mode 100644 index 0000000..5ad1c55 --- /dev/null +++ b/dino-pretraining/tests/test_mil_stream.py @@ -0,0 +1,58 @@ +"""Check replay gradients against a dense bag, including trainable features.""" +import copy + +import torch +from torch import nn + +from mr_dino.mil_head import SparseClassifyThenAggregate +from mr_dino.mil_stream import LastBlockAdapter, backward, evidence +from mr_dino.model import DinoVisionTransformer3D + + +def test_replayed_head_and_feature_gradients_match_dense_global_softmax(): + torch.manual_seed(3) + head=SparseClassifyThenAggregate(16,3,hidden_dim=8,mlp_hidden_dim=7,use_norm=True,checkpoint_chunks=False) + head.route_bias_raw=nn.Parameter(torch.tensor([1.,2.,3.])) + transform=nn.Linear(16,16) + reference=copy.deepcopy(head);reference_transform=copy.deepcopy(transform) + parts=[torch.randn(7,16),torch.randn(3,16),torch.randn(11,16)] + masks=[torch.rand(len(x),3)>.4 for x in parts] + def factory(): + for x,m in zip(parts,masks):yield transform(x),m + target=torch.tensor([1.,0.,1.]) + loss,logits=backward(head,factory,target) + all_tokens=reference_transform(torch.cat(parts)) + scores,patch=reference._chunk_outputs(all_tokens,torch.arange(3)) + scores=scores-(~torch.cat(masks)).float()*torch.nn.functional.softplus(reference.route_bias_raw) + expected=(scores.softmax(0)*patch).sum(0)*reference.output_scale+reference.output_bias + expected_loss=torch.nn.functional.binary_cross_entropy_with_logits(expected,target) + expected_loss.backward() + torch.testing.assert_close(logits[0],expected,atol=1e-6,rtol=1e-5) + torch.testing.assert_close(loss,expected_loss) + for actual,ref in zip(head.parameters(),reference.parameters()): + torch.testing.assert_close(actual.grad,ref.grad,atol=2e-6,rtol=2e-5) + for actual,ref in zip(transform.parameters(),reference_transform.parameters()): + torch.testing.assert_close(actual.grad,ref.grad,atol=2e-6,rtol=2e-5) + expected_evidence=(scores.softmax(0)*patch*reference.output_scale).clamp_min(0) + torch.testing.assert_close(torch.cat(list(evidence(head,factory))),expected_evidence) + + +def test_mri_last_block_matches_backbone_with_dynamic_prefix_and_grid(): + torch.manual_seed(0) + model=DinoVisionTransformer3D(embed_dim=32,depth=2,num_heads=4, + volume_size=(4,32,32),patch_size=(2,16,16), + n_storage_tokens=2,drop_path_rate=.3) + model.init_weights();model.eval() + x=torch.randn(1,1,4,32,32) + tokens,grid=model.prepare_tokens_with_masks(x) + rope=model.rope_embed(D=grid[0],H=grid[1],W=grid[2]) + for block in model.blocks[:-1]:tokens=block(tokens,rope) + model.requires_grad_(False) + adapter=LastBlockAdapter(model) + adapter.train() + actual=adapter(tokens,grid) + expected=model.forward_features(x)["x_norm_patchtokens"] + torch.testing.assert_close(actual,expected) + actual.square().sum().backward() + assert adapter.block.attn.qkv.weight.grad is not None + assert adapter.prefix==3 diff --git a/dino-pretraining/tests/test_fora_recipe.py b/dino-pretraining/tests/test_recipe.py similarity index 98% rename from dino-pretraining/tests/test_fora_recipe.py rename to dino-pretraining/tests/test_recipe.py index 3a4e4e2..b26ec6e 100644 --- a/dino-pretraining/tests/test_fora_recipe.py +++ b/dino-pretraining/tests/test_recipe.py @@ -1,4 +1,4 @@ -"""Tests for the recipe fixes ported from the FORA CT DINOv3 pipeline (mr_dino/recipe.py).""" +"""Tests for the recipe fixes ported from the MR DINOv3 pipeline (mr_dino/recipe.py).""" import math import pytest diff --git a/dino-pretraining/tests/test_vlm_stream.py b/dino-pretraining/tests/test_vlm_stream.py new file mode 100644 index 0000000..d041fdc --- /dev/null +++ b/dino-pretraining/tests/test_vlm_stream.py @@ -0,0 +1,73 @@ +"""Compare the MRI port of MR-RATE with dense autograd, including two ranks.""" +import copy +from datetime import timedelta +from pathlib import Path +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn.functional as F +from mr_dino.vlm_core import StreamPool,pooled_logits_reference,contrastive_ce,streamed_step + + +def test_stream_pool_matches_dense_values_and_every_gradient(): + torch.manual_seed(8) + visual=torch.nn.Linear(7,5);text=torch.nn.Linear(7,5) + visual_ref=copy.deepcopy(visual);text_ref=copy.deepcopy(text) + log_t=torch.tensor(1.2,requires_grad=True);log_p=torch.tensor(.8,requires_grad=True) + ref_t=log_t.detach().clone().requires_grad_();ref_p=log_p.detach().clone().requires_grad_() + image=torch.randn(17,7);words=torch.randn(4,7) + t=F.normalize(text(words),dim=-1);pt=log_p.exp() + pool=StreamPool(t,pt.detach()) + with torch.no_grad(): + for x in image.split(6):pool.add(visual(x)) + p=pool.pooled().requires_grad_(True) + logits=(F.normalize(p,dim=-1)*t).sum(1)*log_t.exp() + loss=logits.square().mean();loss.backward(retain_graph=True) + t_leaf=t.detach().requires_grad_();p_leaf=pt.detach().requires_grad_() + for x in image.split(6):pool.surrogate(t_leaf,p_leaf,visual(x),p.detach(),p.grad).backward() + torch.autograd.backward((t,pt),(t_leaf.grad,p_leaf.grad)) + dense=pooled_logits_reference(F.normalize(text_ref(words),dim=-1),visual_ref(image),ref_p.exp(),ref_t.exp()) + dense.square().mean().backward() + torch.testing.assert_close(logits,dense) + for a,b in zip(list(visual.parameters())+list(text.parameters())+[log_t,log_p], + list(visual_ref.parameters())+list(text_ref.parameters())+[ref_t,ref_p]): + torch.testing.assert_close(a.grad,b.grad,atol=2e-5,rtol=3e-5) + + +def networks(): + torch.manual_seed(17) + return torch.nn.Linear(6,4),torch.nn.Linear(6,4),torch.nn.Parameter(torch.tensor(1.2)),torch.nn.Parameter(torch.tensor(.8)) + + +def data(rank): + generator=torch.Generator().manual_seed(300+rank) + return torch.randn(3+4*rank,6,generator=generator),torch.randn(2,6,generator=generator) + + +def distributed_worker(rank,init_file,result_file): + torch.set_num_threads(1) + dist.init_process_group('gloo',init_method='file://'+init_file,rank=rank,world_size=2,timeout=timedelta(seconds=90)) + try: + visual,text,t,p=networks();image,words=data(rank) + def factory(): + for chunk in image.split(3):yield visual(chunk) + loss,_,_,_=streamed_step(F.normalize(text(words),dim=-1),torch.ones(2,dtype=torch.bool),factory,t,p,2) + params=list(visual.parameters())+list(text.parameters())+[t,p] + flat=torch.cat([x.grad.flatten() for x in params]);dist.all_reduce(flat);flat/=2 + if rank==0:torch.save({'loss':loss,'gradient':flat},result_file) + finally:dist.destroy_process_group() + + +def test_two_rank_full_step_matches_dense_global_objective(tmp_path): + result=tmp_path/'result.pt' + mp.spawn(distributed_worker,args=(str(tmp_path/'rendezvous'),str(result)),nprocs=2,join=True) + actual=torch.load(result,weights_only=True) + visual,text,t,p=networks() + rows=[data(rank) for rank in range(2)] + sentences=F.normalize(text(torch.cat([r[1] for r in rows])),dim=-1) + logits=torch.stack([pooled_logits_reference(sentences,visual(r[0]),p.exp(),t.exp()) for r in rows],dim=1) + loss=contrastive_ce(logits,2,torch.ones(4,dtype=torch.bool),torch.device('cpu')) + loss.backward() + expected=torch.cat([x.grad.flatten() for x in list(visual.parameters())+list(text.parameters())+[t,p]]) + torch.testing.assert_close(actual['loss'],loss) + torch.testing.assert_close(actual['gradient'],expected,atol=2e-5,rtol=3e-5) diff --git a/report-generation-training/README.md b/report-generation-training/README.md index 506fc7e..0bdba41 100644 --- a/report-generation-training/README.md +++ b/report-generation-training/README.md @@ -1,7 +1,7 @@ # MR-RATE single-writer report training This is a standalone report-generation training layer for MR-RATE. It does not -modify FORA or an existing MR-RATE checkout. +modify MR-RATE or an existing MR-RATE checkout. The writer target is the natural `findings` field from `MR-RATE-validation/reports/all_reports.csv`. The standardized @@ -211,8 +211,8 @@ given. ## 8. Evaluation: clinical accuracy and NLG metrics Clinical accuracy re-extracts pathology labels from the generated reports -with the NeuroVFM MRI diagnosis prompt used to build the current 74-label -ground truth (`data-preprocessing/.../07_neurovfm_diagnosis_extraction`), +with the MRI MRI diagnosis prompt used to build the current 74-label +ground truth (`data-preprocessing/.../07_mri_diagnosis_extraction`), then compares them against the ground-truth labels CSV. The bundled schema contains the exact 74 diagnosis keys and guidance used for the completed MR-RATE extraction. On a GPU node with vLLM and Gemma 4 available: @@ -226,8 +226,8 @@ export WORK_DIR="runs/label_extraction_test" sbatch scripts/slurm_extract_labels.sh ``` -The launcher uses the bundled `neurovfm_mri_diagnoses.json` by default. Set -`DIAGNOSES_JSON` only to evaluate a deliberately different NeuroVFM-format +The launcher uses the bundled `mri_diagnoses_74.json` by default. Set +`DIAGNOSES_JSON` only to evaluate a deliberately different MRI-format schema; the vLLM backend rejects the legacy 37-pathology JSON so the two label definitions cannot be mixed silently. diff --git a/report-generation-training/configs/base.yaml b/report-generation-training/configs/base.yaml index 8d56663..81040f7 100644 --- a/report-generation-training/configs/base.yaml +++ b/report-generation-training/configs/base.yaml @@ -11,8 +11,8 @@ data: data_folder: /hnvme/workspace/b180dc51-sezgin/MR-RATE-validation/mri jsonl_file: /hnvme/workspace/b180dc51-sezgin/MR-RATE/contrastive-pretraining/data/findings_sentences.jsonl reports_csv: /hnvme/workspace/b180dc51-sezgin/MR-RATE-validation/reports/all_reports.csv - labels_file: /hnvme/workspace/b180dc51-sezgin/neurovfm_mrrate/splits_neurovfm74/mrrate_neurovfm74_labels.csv - splits_csv: /hnvme/workspace/b180dc51-sezgin/neurovfm_mrrate/splits_neurovfm74/splits.csv + labels_file: /hnvme/workspace/b180dc51-sezgin/mri_mrrate/splits_mri74/mrrate_mri74_labels.csv + splits_csv: /hnvme/workspace/b180dc51-sezgin/mri_mrrate/splits_mri74/splits.csv cached_tokens_dir: artifacts/exact_tokens use_preprocessed: false preprocessed_dir: null diff --git a/report-generation-training/scripts/slurm_extract_labels.sh b/report-generation-training/scripts/slurm_extract_labels.sh index 50410b1..61985a3 100755 --- a/report-generation-training/scripts/slurm_extract_labels.sh +++ b/report-generation-training/scripts/slurm_extract_labels.sh @@ -25,7 +25,7 @@ if [[ ! -d "$project_dir/src/mrrate_report_training" ]]; then echo "Cannot locate report-generation-training; set PROJECT_DIR explicitly" >&2 exit 2 fi -diagnoses_json="${DIAGNOSES_JSON:-$project_dir/../data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_neurovfm_diagnosis_extraction/data/neurovfm_mri_diagnoses.json}" +diagnoses_json="${DIAGNOSES_JSON:-$project_dir/../data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_mri_diagnosis_extraction/data/mri_diagnoses_74.json}" extra_pip="${EXTRA_PIP:-/hnvme/workspace/b180dc51-sezgin/extra-pip}" host_pythonpath="$extra_pip:$project_dir/src${PYTHONPATH:+:$PYTHONPATH}" export PYTHONPATH="$host_pythonpath" diff --git a/report-generation-training/src/mrrate_report_training/extract_labels.py b/report-generation-training/src/mrrate_report_training/extract_labels.py index 23c8382..c598103 100644 --- a/report-generation-training/src/mrrate_report_training/extract_labels.py +++ b/report-generation-training/src/mrrate_report_training/extract_labels.py @@ -1,10 +1,10 @@ -"""Extract NeuroVFM diagnosis labels from generated reports. +"""Extract MRI diagnosis labels from generated reports. -The ``vllm`` backend uses the same NeuroVFM MRI diagnosis extraction prompt, +The ``vllm`` backend uses the same MRI MRI diagnosis extraction prompt, 74-diagnosis guidance schema, parser, and deterministic Gemma settings used to build the current MR-RATE labels. Generated findings are staged as ``batch00_reports.csv`` and passed to -``data-preprocessing/.../07_neurovfm_diagnosis_extraction``. +``data-preprocessing/.../07_mri_diagnosis_extraction``. The ``keyword`` backend is a deterministic name/synonym matcher with basic negation handling. It exists so unit tests and the dummy end-to-end trial can @@ -16,7 +16,7 @@ CLI: python -m mrrate_report_training.extract_labels \ --generated-csv generated_val.csv \ - --diagnoses-json neurovfm_mri_diagnoses.json \ + --diagnoses-json mri_diagnoses_74.json \ --output-csv pred_labels_val.csv \ --backend vllm --work-dir runs/label_extraction_val """ @@ -43,10 +43,10 @@ / "src" / "mr_rate_preprocessing" / "reports_preprocessing" - / "07_neurovfm_diagnosis_extraction" + / "07_mri_diagnosis_extraction" ) _DEFAULT_DIAGNOSES_JSON = ( - _DEFAULT_CLASSIFIER_DIR / "data" / "neurovfm_mri_diagnoses.json" + _DEFAULT_CLASSIFIER_DIR / "data" / "mri_diagnoses_74.json" ) _NEGATIONS = {"no", "not", "without", "absent", "negative", "denies", "denied"} @@ -188,7 +188,7 @@ def _resume_labels_dir( input_hash.update(str(row[text_column] or "").encode()) input_hash.update(b"\0") manifest = { - "format": "mrrate_neurovfm_extraction_v1", + "format": "mrrate_mri_extraction_v1", "studies": len(rows), "input_sha256": input_hash.hexdigest(), "diagnoses_sha256": hashlib.sha256(diagnoses_json.read_bytes()).hexdigest(), @@ -230,19 +230,19 @@ def extract_vllm_labels( batch_size: int, seed: int, ) -> list[dict]: - classifier = classifier_dir / "extract_neurovfm_dx_gemma.py" + classifier = classifier_dir / "extract_mri_dx_gemma.py" merger = classifier_dir / "merge_labels.py" if not classifier.exists() or not merger.exists(): raise FileNotFoundError( - f"NeuroVFM diagnosis extractor not found under {classifier_dir}" + f"MRI diagnosis extractor not found under {classifier_dir}" ) # Refuse the legacy 37-pathology schema: it belongs to the historical - # three-step classifier and does not carry the NeuroVFM diagnosis guidance. + # three-step classifier and does not carry the MRI diagnosis guidance. payload = json.loads(diagnoses_json.read_text()) diagnoses = payload.get("diagnoses") if isinstance(payload, dict) else None if not isinstance(diagnoses, list) or not diagnoses: raise ValueError( - f"{diagnoses_json}: the vllm backend requires the NeuroVFM " + f"{diagnoses_json}: the vllm backend requires the MRI " 'format {"diagnoses": [{"key": ..., "guidance": ...}, ...]}' ) reports_dir, empty = _stage_reports(rows, work_dir, text_column) @@ -251,7 +251,7 @@ def extract_vllm_labels( {"study_uid": row["study_uid"], "labels": {name: 0 for name in schema}} for row in rows ] - # Keep rank files across requeues; the NeuroVFM extractor resumes globally + # Keep rank files across requeues; the MRI extractor resumes globally # by study_uid and checkpoints atomically after every batch. labels_dir = _resume_labels_dir( rows, @@ -331,7 +331,7 @@ def main() -> None: schema_group = parser.add_mutually_exclusive_group() schema_group.add_argument( "--diagnoses-json", - help="NeuroVFM diagnoses JSON; defaults to the bundled 74-diagnosis schema", + help="MRI diagnoses JSON; defaults to the bundled 74-diagnosis schema", ) schema_group.add_argument( "--pathologies-json", diff --git a/report-generation-training/tests/test_extract_labels.py b/report-generation-training/tests/test_extract_labels.py index 0a0225c..6de089f 100644 --- a/report-generation-training/tests/test_extract_labels.py +++ b/report-generation-training/tests/test_extract_labels.py @@ -66,7 +66,7 @@ def test_load_pathology_schema_formats(tmp_path): load_pathology_schema(empty) -def test_load_neurovfm_diagnoses_format(tmp_path): +def test_load_mri_diagnoses_format(tmp_path): diagnoses = tmp_path / "diagnoses.json" diagnoses.write_text( json.dumps( @@ -91,7 +91,7 @@ def test_load_neurovfm_diagnoses_format(tmp_path): assert schema["brain_abscess"][-1] == "cerebral abscess" -def test_bundled_extractor_is_current_neurovfm74_pipeline(): +def test_bundled_extractor_is_current_mri74_pipeline(): schema = load_pathology_schema(_DEFAULT_DIAGNOSES_JSON) assert len(schema) == 74 assert list(schema)[:3] == [ @@ -100,19 +100,19 @@ def test_bundled_extractor_is_current_neurovfm74_pipeline(): "brain_contusion", ] assert list(schema)[-1] == "spine_syrinx" - assert (_DEFAULT_CLASSIFIER_DIR / "extract_neurovfm_dx_gemma.py").is_file() + assert (_DEFAULT_CLASSIFIER_DIR / "extract_mri_dx_gemma.py").is_file() assert (_DEFAULT_CLASSIFIER_DIR / "merge_labels.py").is_file() -def test_bundled_prompt_uses_neurovfm_classification_rules(): - source = (_DEFAULT_CLASSIFIER_DIR / "extract_neurovfm_dx_gemma.py").read_text() +def test_bundled_prompt_uses_mri_classification_rules(): + source = (_DEFAULT_CLASSIFIER_DIR / "extract_mri_dx_gemma.py").read_text() tree = ast.parse(source) assignment = next( node for node in tree.body if isinstance(node, ast.Assign) and any( - isinstance(target, ast.Name) and target.id == "NEUROVFM_PROMPT" + isinstance(target, ast.Name) and target.id == "MRI_DIAGNOSIS_PROMPT" for target in node.targets ) ) @@ -126,7 +126,7 @@ def test_bundled_prompt_uses_neurovfm_classification_rules(): def test_vllm_backend_rejects_legacy_pathology_schema(tmp_path): legacy = tmp_path / "legacy.json" legacy.write_text(json.dumps({"pathologies": {"Gliosis": {}}})) - with pytest.raises(ValueError, match="NeuroVFM"): + with pytest.raises(ValueError, match="MRI"): extract_vllm_labels( [{"study_uid": "s1", "findings_pred": "Gliosis."}], load_pathology_schema(legacy), @@ -140,7 +140,7 @@ def test_vllm_backend_rejects_legacy_pathology_schema(tmp_path): ) -def test_vllm_wrapper_invokes_neurovfm_extractor_and_merges(tmp_path, monkeypatch): +def test_vllm_wrapper_invokes_mri_extractor_and_merges(tmp_path, monkeypatch): diagnoses = tmp_path / "diagnoses.json" diagnoses.write_text( json.dumps( @@ -154,7 +154,7 @@ def test_vllm_wrapper_invokes_neurovfm_extractor_and_merges(tmp_path, monkeypatc ) classifier_dir = tmp_path / "classifier" classifier_dir.mkdir() - (classifier_dir / "extract_neurovfm_dx_gemma.py").write_text("") + (classifier_dir / "extract_mri_dx_gemma.py").write_text("") (classifier_dir / "merge_labels.py").write_text("") calls = [] @@ -186,7 +186,7 @@ def fake_run(command, *, check, env): batch_size=2, seed=42, ) - assert calls[0][1].endswith("extract_neurovfm_dx_gemma.py") + assert calls[0][1].endswith("extract_mri_dx_gemma.py") assert "--diagnoses_json" in calls[0] assert "--pathologies_json" not in calls[0] assert labeled == [