diff --git a/contrastive-pretraining/README.md b/contrastive-pretraining/README.md index ba694d5..368863c 100644 --- a/contrastive-pretraining/README.md +++ b/contrastive-pretraining/README.md @@ -601,6 +601,105 @@ Outputs in `--results_dir`: The best epoch by val mean-AUROC is restored before test evaluation. Single-class columns (all 0 or all 1 in the test split) are gracefully reported as `NaN` and excluded from the macro mean. +## Frozen-encoder MIL probing + +MIL keeps the visual encoder frozen and trains only a gated +Classify-Then-Aggregate head. Unlike linear probing, which uses one pooled vector +per study, MIL operates on the projected visual tokens from every valid series. + +Two execution modes are available: + +| Mode | Token embeddings | Best use | +|------|------------------|----------| +| Cached MIL | Written once as ragged per-study bags | Repeated experiments and faster training | +| Online MIL | Recomputed each epoch and discarded | Limited storage or one-off experiments | + +The encoder configuration must match the contrastive checkpoint, including +`--encoder`, `--fusion_mode`, `--pooling_strategy`, and `--dim_latent`. MIL +requires `--fusion_mode late` so valid series remain separate before their tokens +are concatenated into each study bag. Add `--extra_latent_projection` when that +projection was used during contrastive pretraining. + +### Cached MIL + +Extract token-level features for every split: + +```bash +WEIGHTS=/path/to/model_checkpoint.pt +DATA=/path/to/mri +REPORTS=/path/to/reports.jsonl +LABELS=/path/to/labels.csv +SPLITS=/path/to/splits.csv +FEATURES=/path/to/mil_features + +for SPLIT in train val test; do + python scripts/extract_features.py \ + --weights_path "$WEIGHTS" \ + --encoder vjepa2 \ + --fusion_mode late \ + --pooling_strategy simple_attn \ + --dim_latent 512 \ + --data_folder "$DATA" \ + --jsonl_file "$REPORTS" \ + --labels_file "$LABELS" \ + --splits_csv "$SPLITS" \ + --split "$SPLIT" \ + --normalizer zscore \ + --feature_level tokens \ + --cache_dtype float16 \ + --out_dir "$FEATURES" +done +``` + +Train and evaluate the MIL head: + +```bash +python scripts/mil_probe.py \ + --features_dir "$FEATURES" \ + --results_dir /path/to/mil_results \ + --epochs 50 \ + --batch_size 4 +``` + +Cached MIL stores all retained token embeddings plus study offsets and provenance. +The pooled `features_.npy` files used by the linear probe cannot be reused +for MIL. To bound cache size, add `--max_tokens_per_study 2048` during extraction. +Leaving it at `0` retains every token; any positive limit is deterministic +subsampling. + +### Online MIL without a token cache + +Online MIL reconstructs the same frozen encoder and sends each study's tokens +directly to the same MIL head: + +```bash +python scripts/mil_probe_online.py \ + --weights_path "$WEIGHTS" \ + --encoder vjepa2 \ + --fusion_mode late \ + --data_folder "$DATA" \ + --jsonl_file "$REPORTS" \ + --labels_file "$LABELS" \ + --splits_csv "$SPLITS" \ + --results_dir /path/to/mil_online_results \ + --epochs 50 \ + --grad_accum_steps 4 +``` + +The encoder remains in evaluation mode, runs under `torch.no_grad()`, and is +excluded from the optimizer. Token embeddings are discarded after each study and +are never written to disk. This avoids the token-cache storage cost but repeats +encoder computation every epoch. + +Studies are encoded one at a time because their numbers of series differ. +`--grad_accum_steps` provides a larger effective head-training batch size. +`--max_tokens_per_study` can reduce MIL-head memory after encoding but does not +reduce encoder computation. Preprocessed image inputs remain supported through +`--use_preprocessed --preprocessed_dir /path/to/preprocessed`. + +Both modes select checkpoints and class thresholds using validation data only, +then evaluate the fixed model and thresholds on the test split. + ### Reusing features for a different label set (no re-extract) Because labels are not baked into the encoder features, you only ever run `extract_features.py` **once**. To probe a different label CSV — a different ground-truth rule, or a different grouping — relabel the cached features instead of re-encoding: diff --git a/contrastive-pretraining/mr_rate/mr_rate/mr_rate.py b/contrastive-pretraining/mr_rate/mr_rate/mr_rate.py index c0f347b..008b129 100644 --- a/contrastive-pretraining/mr_rate/mr_rate/mr_rate.py +++ b/contrastive-pretraining/mr_rate/mr_rate/mr_rate.py @@ -340,6 +340,48 @@ def _encode_visual_tokens(self, image, real_volume_mask, vis_proj_layer, return merged, token_mask + def _encode_visual_instances(self, image, real_volume_mask, vis_proj_layer): + """Return every valid series token as one study-level MIL bag. + + Contrastive late fusion averages corresponding tokens across series. + Classify-Then-Aggregate MIL must instead receive the complete set of + series-specific tokens and perform the first study-level aggregation. + """ + if self.fusion_mode != "late": + raise ValueError( + "Series-preserving MIL extraction requires fusion_mode='late'. " + "Other fusion modes aggregate series before projected tokens are available." + ) + + b, r, _c, _d, _h, _w = image.shape + per_series = [] + tokens_per_series = None + for series_index in range(r): + rank = dist.get_rank() if dist.is_initialized() else 0 + if rank == 0: + print( + f"[MRRATE rank=0] encoding MIL series {series_index + 1}/{r}...", + flush=True, + ) + encoded = self.run_checkpoint(self.visual_transformer, image[:, series_index]) + projected = vis_proj_layer(encoded) + if projected.ndim != 3: + raise RuntimeError( + f"Visual encoder must return [B,T,D] tokens, got {projected.shape}" + ) + if tokens_per_series is None: + tokens_per_series = projected.shape[1] + elif projected.shape[1] != tokens_per_series: + raise RuntimeError("All padded series must produce the same token count") + per_series.append(projected) + + stacked = torch.stack(per_series, dim=1) # [B,R,T,D] + token_mask = real_volume_mask.unsqueeze(-1).expand(b, r, tokens_per_series) + return ( + stacked.reshape(b, r * tokens_per_series, stacked.shape[-1]), + token_mask.reshape(b, r * tokens_per_series), + ) + def state_dict(self, *args, **kwargs): return super().state_dict(*args, **kwargs) @@ -370,6 +412,7 @@ def forward( aug_text=None, mode='clip', debug=False, + return_visual_tokens=False, **kwargs ): b, r, c, d, h, w = image.shape @@ -380,7 +423,10 @@ def forward( if real_volume_mask is None: real_volume_mask = torch.ones(b, r, dtype=torch.bool, device=device) - use_extra_proj = self.extra_latent_projection and return_loss + if return_visual_tokens and return_loss: + raise ValueError("return_visual_tokens and return_loss are mutually exclusive") + + use_extra_proj = self.extra_latent_projection and (return_loss or return_visual_tokens) vis_proj_layer = self.to_visual_latent_extra if use_extra_proj else self.to_visual_latent text_proj_layer = self.to_text_latent_extra if use_extra_proj else self.to_text_latent @@ -406,15 +452,26 @@ def forward( print(f"[DEBUG MRRATE rank={dist.get_rank()}] text done, encoding visual...", flush=True) # 2. Encode Visual Tokens (Returns Tokens + Mask) - visual_tokens, token_mask = self._encode_visual_tokens( - image, real_volume_mask, vis_proj_layer, - text_latents=text_latents, - num_sentences_per_image=num_sentences_per_image - ) + if return_visual_tokens: + visual_tokens, token_mask = self._encode_visual_instances( + image, real_volume_mask, vis_proj_layer + ) + else: + visual_tokens, token_mask = self._encode_visual_tokens( + image, real_volume_mask, vis_proj_layer, + text_latents=text_latents, + num_sentences_per_image=num_sentences_per_image, + ) if debug and dist.is_initialized(): print(f"[DEBUG MRRATE rank={dist.get_rank()}] visual done: tokens={visual_tokens.shape}, mask={token_mask.shape}", flush=True) + # Downstream MIL consumes the projected visual instances before the + # global masked mean used by standard MR-RATE inference. This opt-in + # branch leaves contrastive training and existing inference unchanged. + if return_visual_tokens: + return visual_tokens, token_mask + # 3. Training Loss if return_loss: temp = self.logit_temperature.float().clamp(min=0.0, max=4.0).exp() @@ -518,4 +575,4 @@ def forward( mask_bc = token_mask.unsqueeze(-1) sum_tokens = (visual_tokens * mask_bc).sum(dim=1) count_tokens = mask_bc.sum(dim=1).clamp(min=1.0) - return l2norm(sum_tokens / count_tokens) \ No newline at end of file + return l2norm(sum_tokens / count_tokens) diff --git a/contrastive-pretraining/scripts/extract_features.py b/contrastive-pretraining/scripts/extract_features.py index eeeb830..3a09a25 100644 --- a/contrastive-pretraining/scripts/extract_features.py +++ b/contrastive-pretraining/scripts/extract_features.py @@ -1,21 +1,28 @@ """ -Extract frozen MR-RATE visual features for linear probing. +Extract frozen MR-RATE visual features for linear or MIL probing. Given a pretrained MR-RATE checkpoint, runs the visual encoder + projection + masked pooling over every subject in a split and dumps: - /features_.npy float32 [N, dim_latent] - /labels_.npy float32 [N, num_classes] - /subject_ids_.txt one study_uid per line - /label_names.json list of 32 pathology names - (only written on the first run) +Pooled mode (default) writes the existing linear-probe contract: + + /features_.npy float32 [N, dim_latent] + +Token mode writes a ragged, memory-mappable MIL contract instead: + + /tokens_.bin concatenated valid projected tokens + /token_offsets_.npy int64 [N + 1] bag boundaries + /token_features_.json cache metadata and provenance + +Both modes also write labels, study IDs, and the ordered label schema. Run this once per split (train / val / test). Then `linear_probe.py` trains and evaluates a linear classifier on the cached features in seconds — no need to re-encode the 3D volumes every epoch. -The frozen backbone returns `l2norm(masked_mean(visual_tokens))`, the -same global representation that `inference.py` uses for zero-shot scoring. +Token mode stops immediately before MR-RATE's final masked mean. It is the +required input for NeuroVFM-style Classify-Then-Aggregate MIL; pooled +`features_*.npy` cannot be used for MIL because the instances are gone. Usage: python extract_features.py \ @@ -33,6 +40,8 @@ import os import json import argparse +import hashlib +import uuid from pathlib import Path import numpy as np @@ -49,6 +58,80 @@ torch.backends.cudnn.allow_tf32 = True +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _file_record(path: str, include_digest: bool = True) -> dict: + resolved = Path(path).resolve() + stat = resolved.stat() + record = { + "path": str(resolved), + "size": stat.st_size, + "mtime_ns": stat.st_mtime_ns, + } + if include_digest: + record["sha256"] = _sha256_file(resolved) + return record + + +def _atomic_save_npy(path: Path, array: np.ndarray) -> None: + temporary = path.with_name(path.name + f".tmp.{os.getpid()}") + with temporary.open("wb") as handle: + np.save(handle, array) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + + +def _atomic_write_text(path: Path, value: str) -> None: + temporary = path.with_name(path.name + f".tmp.{os.getpid()}") + with temporary.open("w") as handle: + handle.write(value) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + + +def _cache_provenance(args, dim_latent: int) -> tuple[dict, str]: + jsonl = _file_record(args.jsonl_file, include_digest=False) + provenance = { + "version": 2, + "checkpoint": _file_record(args.weights_path), + "encoder": args.encoder, + "vjepa21_checkpoint": ( + _file_record(args.vjepa21_checkpoint) + if args.vjepa21_checkpoint else None + ), + "chunk_size": args.chunk_size, + "fusion_mode": args.fusion_mode, + "pooling_strategy": args.pooling_strategy, + "extra_latent_projection": args.extra_latent_projection, + "dim_latent": dim_latent, + "normalizer": args.normalizer, + "space": args.space, + "use_preprocessed": args.use_preprocessed, + "preprocessed_dir": ( + str(Path(args.preprocessed_dir).resolve()) + if args.preprocessed_dir else None + ), + "data_folder": ( + str(Path(args.data_folder).resolve()) if args.data_folder else None + ), + "jsonl": jsonl, + "labels": _file_record(args.labels_file), + "splits": _file_record(args.splits_csv), + "cache_dtype": args.cache_dtype, + "max_tokens_per_study": args.max_tokens_per_study, + } + serialized = json.dumps(provenance, sort_keys=True, separators=(",", ":")) + return provenance, hashlib.sha256(serialized.encode("utf-8")).hexdigest() + + def _load_and_verify(clip: "MRRATE", weights_path: str, strict_missing: bool = False) -> None: """Load checkpoint into MRRATE and report exactly what was matched. @@ -61,7 +144,6 @@ def _load_and_verify(clip: "MRRATE", weights_path: str, strict_missing: bool = F 3) Optionally aborts on any missing key (--strict_missing) so a typo in --fusion_mode (which changes module names) fails loudly. """ - import hashlib import torch as _torch from pathlib import Path as _Path @@ -72,12 +154,27 @@ def _load_and_verify(clip: "MRRATE", weights_path: str, strict_missing: bool = F # Snapshot a representative set of param hashes pre-load def _hashes(model): out = {} + group_counts = {"projection": 0, "pool": 0, "encoder": 0, "text": 0} for n, t in model.named_parameters(): - # Sample only a handful for speed; covers projections, pooling, encoder - if any(s in n for s in ("to_visual_latent", "to_text_latent", - "recon_pool", "visual_transformer", - "text_transformer.encoder.layer.0")): - out[n] = hashlib.md5(t.detach().cpu().float().numpy().tobytes()).hexdigest()[:8] + if "to_visual_latent" in n: + group = "projection" + elif "recon_pool" in n: + group = "pool" + elif "visual_transformer" in n: + group = "encoder" + elif "text_transformer.encoder.layer.0" in n: + group = "text" + else: + continue + if group_counts[group] >= 2: + continue + flat = t.detach().reshape(-1) + if flat.numel() > 4096: + indices = _torch.linspace(0, flat.numel() - 1, 4096).long() + flat = flat.index_select(0, indices.to(flat.device)) + value = flat.cpu().float().numpy().tobytes() + out[n] = hashlib.md5(value).hexdigest()[:8] + group_counts[group] += 1 return out pre = _hashes(clip) @@ -108,8 +205,20 @@ def _hashes(model): if changed == 0: raise RuntimeError( f"No model parameters changed when loading {weights_path}. " - "The checkpoint likely doesn't match the architecture " - "(wrong --fusion_mode / --encoder / --dim_latent)." + "The checkpoint does not match the constructed model parameters." + ) + projector_prefix = ( + "to_visual_latent_extra." + if clip.extra_latent_projection else "to_visual_latent." + ) + critical_prefixes = ("visual_transformer.", projector_prefix) + critical_missing = [ + name for name in missing if name.startswith(critical_prefixes) + ] + if critical_missing: + raise RuntimeError( + "Checkpoint is missing MIL-critical visual parameters. " + f"First entries: {critical_missing[:10]}" ) if strict_missing and missing: raise RuntimeError( @@ -166,6 +275,7 @@ def build_encoder(args) -> tuple[MRRATE, int]: dim_latent=args.dim_latent, fusion_mode=args.fusion_mode, pooling_strategy=args.pooling_strategy, + extra_latent_projection=args.extra_latent_projection, use_gradient_checkpointing=False, ).cuda() return clip, args.dim_latent @@ -184,6 +294,8 @@ def main() -> None: parser.add_argument("--pooling_strategy", type=str, default="simple_attn", choices=["simple_attn", "cross_attn", "gated"]) parser.add_argument("--dim_latent", type=int, default=512) + parser.add_argument("--extra_latent_projection", action="store_true", + help="Use the checkpoint's extra visual projection for MIL tokens.") # Data parser.add_argument("--data_folder", type=str, default=None, help="Raw MR data folder. Required unless --use_preprocessed.") @@ -204,6 +316,16 @@ def main() -> None: help="Downgrade a cache-manifest config mismatch to a warning.") # Output parser.add_argument("--out_dir", type=str, default="./linear_probe_features") + parser.add_argument("--feature_level", type=str, default="pooled", + choices=["pooled", "tokens"], + help="pooled keeps the linear-probe format; tokens writes " + "ragged pre-global-pooling bags for mil_probe.py.") + parser.add_argument("--cache_dtype", type=str, default="float16", + choices=["float16", "float32"], + help="On-disk dtype for --feature_level tokens.") + parser.add_argument("--max_tokens_per_study", type=int, default=0, + help="Deterministically subsample each token bag to this size. " + "0 preserves every encoder token (exact but potentially very large).") parser.add_argument("--strict_missing", action="store_true", help="Abort if the checkpoint is missing any model parameter " "(catches wrong --fusion_mode / --encoder mismatch).") @@ -214,6 +336,13 @@ def main() -> None: parser.error("--use_preprocessed requires --preprocessed_dir") elif not args.data_folder: parser.error("--data_folder is required unless --use_preprocessed is set") + if args.max_tokens_per_study < 0: + parser.error("--max_tokens_per_study must be >= 0") + if args.feature_level == "tokens" and args.fusion_mode != "late": + parser.error( + "MIL token extraction requires --fusion_mode late so every valid " + "series can remain a separate set of instances." + ) out_dir = Path(args.out_dir) out_dir.mkdir(parents=True, exist_ok=True) @@ -257,8 +386,18 @@ def main() -> None: # Persist label names so the linear-probe trainer doesn't need the source CSV names_path = out_dir / "label_names.json" - if not names_path.exists(): - names_path.write_text(json.dumps(ds.label_columns, indent=2, ensure_ascii=False) + "\n") + if names_path.exists(): + existing_names = json.loads(names_path.read_text()) + if existing_names != ds.label_columns: + raise RuntimeError( + f"{names_path} does not match the labels file for this run. " + "Use a separate --out_dir or remove the stale cache." + ) + else: + _atomic_write_text( + names_path, + json.dumps(ds.label_columns, indent=2, ensure_ascii=False) + "\n", + ) print(f"Wrote {names_path}") loader = DataLoader( @@ -266,54 +405,167 @@ def main() -> None: drop_last=False, collate_fn=collate_fn_infer, pin_memory=True, ) - feats: list[np.ndarray] = [] + pooled_feats: list[np.ndarray] = [] labs: list[np.ndarray] = [] sids: list[str] = [] n_unlabeled = 0 device = next(clip.parameters()).device - print(f"\n--- Encoding {len(loader)} subjects ---") - with torch.no_grad(): - for batch in tqdm.tqdm(loader, desc=f"encode[{args.split}]"): - imgs, _sentences, subject_id, real_volume_mask, labels = batch - if labels.size == 0: - n_unlabeled += 1 - continue - imgs = imgs.to(device, dtype=torch.bfloat16) - real_volume_mask = real_volume_mask.to(device) - - with autocast(dtype=torch.bfloat16): - # MRRATE in inference mode returns the masked-mean pooled, - # L2-normalized latent: [1, dim_latent] - pooled = clip( - text_input=None, - image=imgs, - device=device, - real_volume_mask=real_volume_mask, - return_loss=False, + token_dtype = np.dtype(args.cache_dtype) + token_offsets: list[int] = [0] + full_token_counts: list[int] = [] + series_counts: list[int] = [] + cache_provenance = cache_fingerprint = None + token_path = token_tmp_path = None + token_fh = None + cache_id = None + if args.feature_level == "tokens": + cache_provenance, cache_fingerprint = _cache_provenance(args, dim_latent) + for other_manifest in out_dir.glob("token_features_*.json"): + other = json.loads(other_manifest.read_text()) + if other.get("cache_fingerprint") != cache_fingerprint: + raise RuntimeError( + f"Cache configuration differs from {other_manifest}. " + "Use a separate --out_dir for a different encoder or preprocessing setup." ) + cache_id = uuid.uuid4().hex[:12] + token_path = out_dir / f"tokens_{args.split}_{cache_id}.bin" + token_tmp_path = token_path.with_name(token_path.name + f".tmp.{os.getpid()}") + token_fh = open(token_tmp_path, "wb") - feats.append(pooled.float().cpu().numpy().reshape(-1)) - labs.append(np.asarray(labels, dtype=np.float32).reshape(-1)) - sids.append(subject_id) - - if not feats: - raise RuntimeError(f"No labeled subjects encoded for split={args.split}.") + print(f"\n--- Encoding {len(loader)} subjects ---") + try: + with torch.no_grad(): + for batch in tqdm.tqdm(loader, desc=f"encode[{args.split}]"): + imgs, _sentences, subject_id, real_volume_mask, labels = batch + if labels.size == 0: + n_unlabeled += 1 + continue + imgs = imgs.to(device, dtype=torch.bfloat16) + real_volume_mask = real_volume_mask.to(device) + + with autocast(dtype=torch.bfloat16): + encoded = clip( + text_input=None, + image=imgs, + device=device, + real_volume_mask=real_volume_mask, + return_loss=False, + return_visual_tokens=(args.feature_level == "tokens"), + ) + + if args.feature_level == "tokens": + visual_tokens, token_mask = encoded + valid_tokens = visual_tokens[0, token_mask[0]] + if valid_tokens.ndim != 2 or valid_tokens.shape[1] != dim_latent: + raise RuntimeError( + f"Unexpected visual token shape for {subject_id}: " + f"{tuple(valid_tokens.shape)}; expected [T, {dim_latent}]" + ) + if valid_tokens.shape[0] == 0: + raise RuntimeError(f"Encoder returned an empty token bag for {subject_id}") + padded_series = real_volume_mask.shape[1] + if visual_tokens.shape[1] % padded_series != 0: + raise RuntimeError("Flat MIL tokens cannot be divided into series") + tokens_per_series = visual_tokens.shape[1] // padded_series + real_series = int(real_volume_mask[0].sum().item()) + full_token_count = real_series * tokens_per_series + if valid_tokens.shape[0] != full_token_count: + raise RuntimeError("Series mask and valid token count disagree") + full_token_counts.append(full_token_count) + series_counts.append(real_series) + if (args.max_tokens_per_study > 0 and + valid_tokens.shape[0] > args.max_tokens_per_study): + keep_array = np.rint(np.linspace( + 0, valid_tokens.shape[0] - 1, + num=args.max_tokens_per_study, + )).astype(np.int64) + keep = torch.from_numpy(keep_array).to(valid_tokens.device) + valid_tokens = valid_tokens.index_select(0, keep) + token_array = ( + valid_tokens.float().cpu().numpy().astype(token_dtype, copy=False) + ) + token_array.tofile(token_fh) + token_offsets.append(token_offsets[-1] + token_array.shape[0]) + else: + pooled_feats.append(encoded.float().cpu().numpy().reshape(-1)) + + labs.append(np.asarray(labels, dtype=np.float32).reshape(-1)) + sids.append(subject_id) + + if not labs: + raise RuntimeError(f"No labeled subjects encoded for split={args.split}.") + except BaseException: + if token_fh is not None: + token_fh.close() + token_tmp_path.unlink(missing_ok=True) + raise + + if token_fh is not None: + token_fh.flush() + os.fsync(token_fh.fileno()) + token_fh.close() - F = np.stack(feats, axis=0) # [N, dim_latent] Y = np.stack(labs, axis=0) # [N, num_classes] - assert F.shape[0] == Y.shape[0] == len(sids) + assert Y.shape[0] == len(sids) assert Y.shape[1] == num_classes, f"label width {Y.shape[1]} != classes {num_classes}" - feat_path = out_dir / f"features_{args.split}.npy" - lab_path = out_dir / f"labels_{args.split}.npy" - sid_path = out_dir / f"subject_ids_{args.split}.txt" - np.save(feat_path, F) - np.save(lab_path, Y) - sid_path.write_text("\n".join(sids) + "\n") - print(f"\nWrote:") - print(f" {feat_path} shape={F.shape} dtype={F.dtype}") + if args.feature_level == "tokens": + offsets = np.asarray(token_offsets, dtype=np.int64) + if offsets.shape[0] != Y.shape[0] + 1: + raise RuntimeError("Token offsets and labels became misaligned") + os.replace(token_tmp_path, token_path) + offsets_path = out_dir / f"token_offsets_{args.split}_{cache_id}.npy" + lab_path = out_dir / f"labels_{args.split}_{cache_id}.npy" + sid_path = out_dir / f"subject_ids_{args.split}_{cache_id}.txt" + full_counts_path = out_dir / f"full_token_counts_{args.split}_{cache_id}.npy" + series_counts_path = out_dir / f"series_counts_{args.split}_{cache_id}.npy" + metadata_path = out_dir / f"token_features_{args.split}.json" + _atomic_save_npy(offsets_path, offsets) + _atomic_save_npy(lab_path, Y) + _atomic_write_text(sid_path, "\n".join(sids) + "\n") + _atomic_save_npy( + full_counts_path, np.asarray(full_token_counts, dtype=np.int64) + ) + _atomic_save_npy( + series_counts_path, np.asarray(series_counts, dtype=np.int32) + ) + metadata = { + "format": "raw_numpy_memmap", + "format_version": 2, + "feature_level": "projected_per_series_visual_tokens", + "split": args.split, + "tokens_file": token_path.name, + "offsets_file": offsets_path.name, + "labels_file": lab_path.name, + "subject_ids_file": sid_path.name, + "full_token_counts_file": full_counts_path.name, + "series_counts_file": series_counts_path.name, + "dtype": token_dtype.name, + "dim": dim_latent, + "num_studies": len(sids), + "num_tokens": int(offsets[-1]), + "max_tokens_per_study": args.max_tokens_per_study, + "cache_fingerprint": cache_fingerprint, + "provenance": cache_provenance, + } + # Publish the canonical manifest last. Until this atomic replace, any + # previous cache generation remains internally consistent. + _atomic_write_text(metadata_path, json.dumps(metadata, indent=2) + "\n") + print(f" {token_path} shape=({offsets[-1]}, {dim_latent}) dtype={token_dtype.name}") + print(f" {offsets_path} shape={offsets.shape} dtype={offsets.dtype}") + print(f" {metadata_path}") + else: + feature_matrix = np.stack(pooled_feats, axis=0) + assert feature_matrix.shape[0] == Y.shape[0] + feat_path = out_dir / f"features_{args.split}.npy" + lab_path = out_dir / f"labels_{args.split}.npy" + sid_path = out_dir / f"subject_ids_{args.split}.txt" + _atomic_save_npy(feat_path, feature_matrix) + _atomic_save_npy(lab_path, Y) + _atomic_write_text(sid_path, "\n".join(sids) + "\n") + print(f" {feat_path} shape={feature_matrix.shape} dtype={feature_matrix.dtype}") print(f" {lab_path} shape={Y.shape} dtype={Y.dtype}") print(f" {sid_path} ({len(sids)} ids)") if n_unlabeled: diff --git a/contrastive-pretraining/scripts/mil_probe.py b/contrastive-pretraining/scripts/mil_probe.py new file mode 100644 index 0000000..419a11e --- /dev/null +++ b/contrastive-pretraining/scripts/mil_probe.py @@ -0,0 +1,700 @@ +""" +NeuroVFM Classify-Then-Aggregate MIL on frozen MR-RATE token features. + +The head matches NeuroVFM's diagnostic architecture while replacing only +its visual encoder output with projected tokens from a contrastively +pretrained MR-RATE checkpoint: + + attention = W(tanh(V(x)) * sigmoid(U(x))) + patch_logits = MLP(x) + bag_logit[c] = sum_i softmax_i(attention[i, c]) * patch_logits[i, c] + +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 +parameterization and computation are otherwise the same. + +Workflow: + python scripts/extract_features.py ... --split train \ + --feature_level tokens --cache_dtype float16 --out_dir mil_features + python scripts/extract_features.py ... --split val \ + --feature_level tokens --cache_dtype float16 --out_dir mil_features + python scripts/extract_features.py ... --split test \ + --feature_level tokens --cache_dtype float16 --out_dir mil_features + python scripts/mil_probe.py \ + --features_dir mil_features --results_dir mil_results + +Exact token caches can be very large. extract_features.py offers an +explicit --max_tokens_per_study approximation; leaving it at 0 preserves +the complete encoder token bag. +""" +from __future__ import annotations + +import argparse +import csv +import contextlib +import json +import math +import re +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import DataLoader, Dataset + +class PatchClassifier(nn.Module): + """NeuroVFM's shared MLP: Linear/GELU hidden blocks, then Linear.""" + + def __init__(self, dim: int, hidden_dims: tuple[int, ...], out_dim: int): + super().__init__() + layers: list[nn.Module] = [] + current = dim + for hidden in hidden_dims: + layers.extend((nn.Linear(current, hidden), nn.GELU())) + current = hidden + layers.append(nn.Linear(current, out_dim)) + self.net = nn.Sequential(*layers) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +class ClassifyThenAggregate(nn.Module): + """NeuroVFM's 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. + """ + + def __init__( + self, + dim: int, + n_classes: int, + hidden_dim: int = 512, + mlp_hidden_dims: tuple[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, + ): + super().__init__() + self.dim = dim + self.n_classes = n_classes + self.use_gating = use_gating + self.use_output_bias_scale = use_output_bias_scale + self.init_std = init_std + + norm = lambda: nn.LayerNorm(dim, eps=1e-6) if use_norm else nn.Identity() + self.attention_norm = norm() + self.classifier_norm = norm() + self.input_dropout = nn.Dropout(drop_rate) + + self.attention_v = nn.Linear(dim, hidden_dim) + self.attention_u = nn.Linear(dim, hidden_dim) if use_gating else None + self.attention_w = nn.Linear(hidden_dim, n_classes) + self.patch_classifier = PatchClassifier(dim, mlp_hidden_dims, n_classes) + + if use_output_bias_scale: + self.output_scale = nn.Parameter(torch.ones(n_classes)) + self.output_bias = nn.Parameter(torch.zeros(n_classes)) + else: + self.register_parameter("output_scale", None) + self.register_parameter("output_bias", None) + + 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=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 forward( + self, + tokens: torch.Tensor, + cu_seqlens: torch.Tensor, + return_details: bool = False, + ): + if tokens.ndim != 2 or tokens.shape[1] != self.dim: + raise ValueError(f"tokens must have shape [N, {self.dim}], got {tuple(tokens.shape)}") + bounds = cu_seqlens.detach().cpu().tolist() + if len(bounds) < 2 or bounds[0] != 0 or bounds[-1] != tokens.shape[0]: + raise ValueError("cu_seqlens must start at 0 and end at the flat token count") + + attention_input = self.input_dropout(self.attention_norm(tokens)) + classifier_input = self.input_dropout(self.classifier_norm(tokens)) + + attention_hidden = torch.tanh(self.attention_v(attention_input)) + if self.use_gating: + 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: list[torch.Tensor] = [] + attention_chunks: list[torch.Tensor] = [] + for start, end in zip(bounds[:-1], bounds[1:]): + if end <= start: + raise ValueError("MIL bags must contain at least one token") + 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.use_output_bias_scale: + output = output * self.output_scale + self.output_bias + + if return_details: + return output, torch.cat(attention_chunks, dim=0), patch_logits + return output + + +class RaggedTokenDataset(Dataset): + """Memory-mapped projected MR-RATE token bags.""" + + def __init__(self, features_dir: Path, split: str): + self.features_dir = features_dir + self.split = split + metadata_path = features_dir / f"token_features_{split}.json" + if not metadata_path.exists(): + raise FileNotFoundError( + f"Missing {metadata_path}. Run extract_features.py with " + f"--feature_level tokens --split {split}." + ) + self.metadata = json.loads(metadata_path.read_text()) + if self.metadata.get("format") != "raw_numpy_memmap": + raise ValueError(f"Unsupported token cache format in {metadata_path}") + + self.dim = int(self.metadata["dim"]) + self.dtype = np.dtype(self.metadata["dtype"]) + if self.dtype not in (np.dtype("float16"), np.dtype("float32")): + raise ValueError(f"Unsupported token dtype {self.dtype} for split={split}") + self.offsets = np.load(features_dir / self.metadata["offsets_file"]) + labels_file = self.metadata.get("labels_file", f"labels_{split}.npy") + ids_file = self.metadata.get("subject_ids_file", f"subject_ids_{split}.txt") + self.labels = np.load(features_dir / labels_file).astype(np.float32) + ids_path = features_dir / ids_file + self.subject_ids = ids_path.read_text().strip().splitlines() + + full_counts_file = self.metadata.get("full_token_counts_file") + series_counts_file = self.metadata.get("series_counts_file") + if full_counts_file and series_counts_file: + self.full_token_counts = np.load(features_dir / full_counts_file) + self.series_counts = np.load(features_dir / series_counts_file) + else: + self.full_token_counts = np.diff(self.offsets).astype(np.int64) + self.series_counts = np.ones(len(self.labels), dtype=np.int32) + + if self.offsets.ndim != 1 or self.offsets.shape[0] != self.labels.shape[0] + 1: + raise ValueError(f"Offsets/labels mismatch for split={split}") + if len(self.labels) == 0: + raise ValueError(f"Split {split} contains no studies") + if len(self.subject_ids) != self.labels.shape[0]: + raise ValueError(f"Study IDs/labels mismatch for split={split}") + if len(set(self.subject_ids)) != len(self.subject_ids): + raise ValueError(f"Duplicate study IDs within split={split}") + if self.offsets[0] != 0 or np.any(np.diff(self.offsets) <= 0): + raise ValueError(f"Invalid or empty token bags for split={split}") + if not np.isfinite(self.labels).all(): + raise ValueError(f"Non-finite labels in split={split}") + if not np.isin(self.labels, (0.0, 1.0)).all(): + raise ValueError(f"Labels must be strictly binary 0/1 in split={split}") + if not ( + len(self.full_token_counts) == len(self.series_counts) == len(self.labels) + ): + raise ValueError(f"Token mapping metadata mismatch for split={split}") + if np.any(self.series_counts <= 0) or np.any(self.full_token_counts <= 0): + raise ValueError(f"Invalid series/token counts for split={split}") + if np.any(self.full_token_counts % self.series_counts != 0): + raise ValueError(f"Token counts are not divisible by series counts for split={split}") + + self.num_tokens = int(self.offsets[-1]) + token_path = features_dir / self.metadata["tokens_file"] + expected_bytes = self.num_tokens * self.dim * self.dtype.itemsize + if token_path.stat().st_size != expected_bytes: + raise ValueError( + f"Token cache size mismatch for {token_path}: expected " + f"{expected_bytes} bytes, found {token_path.stat().st_size}" + ) + self.tokens = np.memmap( + token_path, + mode="r", + dtype=self.dtype, + shape=(self.num_tokens, self.dim), + ) + + def __len__(self) -> int: + return self.labels.shape[0] + + def __getitem__(self, index: int): + start, end = int(self.offsets[index]), int(self.offsets[index + 1]) + # Copy makes the NumPy buffer writable for safe torch collation. + tokens = torch.from_numpy(np.array(self.tokens[start:end], copy=True)) + labels = torch.from_numpy(self.labels[index]) + return tokens, labels, index + + def token_mapping(self, index: int) -> tuple[np.ndarray, int, int]: + """Return original flat indices, series count, and tokens per series.""" + cached_count = int(self.offsets[index + 1] - self.offsets[index]) + full_count = int(self.full_token_counts[index]) + series_count = int(self.series_counts[index]) + if cached_count == full_count: + token_indices = np.arange(full_count, dtype=np.int64) + else: + token_indices = np.rint( + np.linspace(0, full_count - 1, cached_count) + ).astype(np.int64) + return token_indices, series_count, full_count // series_count + + +def collate_ragged(batch): + token_bags, labels, indices = zip(*batch) + lengths = torch.tensor([bag.shape[0] for bag in token_bags], dtype=torch.long) + cu_seqlens = torch.cat((torch.zeros(1, dtype=torch.long), lengths.cumsum(0))) + return ( + torch.cat(token_bags, dim=0), + torch.stack(labels, dim=0), + cu_seqlens, + torch.tensor(indices, dtype=torch.long), + ) + + +def amp_context(device: torch.device): + if device.type == "cuda" and torch.cuda.is_bf16_supported(): + return torch.autocast(device_type="cuda", dtype=torch.bfloat16) + return contextlib.nullcontext() + + +def prepare_tokens(tokens: torch.Tensor, device: torch.device) -> torch.Tensor: + tokens = tokens.to(device, non_blocking=True) + if device.type != "cuda" or not torch.cuda.is_bf16_supported(): + tokens = tokens.float() + return tokens + + +def auroc_table( + logits: np.ndarray, + labels: np.ndarray, + label_names: list[str], +) -> tuple[float, dict[str, float]]: + from sklearn.metrics import roc_auc_score + + per_class: dict[str, float] = {} + valid: list[float] = [] + for class_index, name in enumerate(label_names): + target = labels[:, class_index] + if len(np.unique(target)) < 2: + per_class[name] = float("nan") + continue + score = float(roc_auc_score(target, logits[:, class_index])) + per_class[name] = score + valid.append(score) + return (float(np.mean(valid)) if valid else float("nan")), per_class + + +def select_validation_thresholds( + logits: np.ndarray, + labels: np.ndarray, +) -> np.ndarray: + """Select one logit threshold per class using validation labels only.""" + from sklearn.metrics import roc_curve + + thresholds = np.zeros(logits.shape[1], dtype=np.float32) + for class_index in range(logits.shape[1]): + target = labels[:, class_index] + if len(np.unique(target)) < 2: + continue + false_positive, true_positive, candidates = roc_curve( + target, logits[:, class_index] + ) + objective = true_positive - false_positive + objective[~np.isfinite(candidates)] = -np.inf + best = int(np.argmax(objective)) + if np.isfinite(candidates[best]): + thresholds[class_index] = float(candidates[best]) + return thresholds + + +def per_class_metrics( + logits: np.ndarray, + labels: np.ndarray, + label_names: list[str], + thresholds: np.ndarray, +) -> list[dict]: + from sklearn.metrics import average_precision_score, roc_auc_score + + rows = [] + for class_index, name in enumerate(label_names): + target = labels[:, class_index] + prediction = logits[:, class_index] >= thresholds[class_index] + positive = target == 1 + negative = ~positive + true_positive = int(np.sum(prediction & positive)) + false_negative = int(np.sum(~prediction & positive)) + true_negative = int(np.sum(~prediction & negative)) + false_positive = int(np.sum(prediction & negative)) + has_both = len(np.unique(target)) == 2 + rows.append({ + "label": name, + "auroc": float(roc_auc_score(target, logits[:, class_index])) if has_both else None, + "auprc": float(average_precision_score(target, logits[:, class_index])) if positive.any() else None, + "threshold": float(thresholds[class_index]), + "sensitivity": ( + true_positive / (true_positive + false_negative) + if true_positive + false_negative else None + ), + "specificity": ( + true_negative / (true_negative + false_positive) + if true_negative + false_positive else None + ), + "positives": int(positive.sum()), + "negatives": int(negative.sum()), + }) + return rows + + +def evaluate_head( + head: ClassifyThenAggregate, + loader: DataLoader, + device: torch.device, + attention_dir: Path | None = None, +) -> tuple[np.ndarray, np.ndarray, list[int]]: + head.eval() + logits_out: list[np.ndarray] = [] + labels_out: list[np.ndarray] = [] + indices_out: list[int] = [] + if attention_dir is not None: + if attention_dir.exists() and any(attention_dir.iterdir()): + raise RuntimeError( + f"Attention directory is not empty: {attention_dir}. " + "Use a new --results_dir to avoid stale artifacts." + ) + attention_dir.mkdir(parents=True, exist_ok=True) + + with torch.no_grad(): + for tokens, labels, cu_seqlens, indices in loader: + tokens = prepare_tokens(tokens, device) + labels = labels.to(device, non_blocking=True) + with amp_context(device): + if attention_dir is None: + logits = head(tokens, cu_seqlens) + attention = patch_logits = None + else: + logits, attention, patch_logits = head( + tokens, cu_seqlens, return_details=True + ) + + logits_out.append(logits.float().cpu().numpy()) + labels_out.append(labels.float().cpu().numpy()) + batch_indices = indices.tolist() + indices_out.extend(batch_indices) + + if attention_dir is not None: + bounds = cu_seqlens.tolist() + dataset = loader.dataset + for local_index, dataset_index in enumerate(batch_indices): + start, end = bounds[local_index], bounds[local_index + 1] + subject_id = dataset.subject_ids[dataset_index] + safe_id = re.sub(r"[^A-Za-z0-9_.-]", "_", subject_id) + token_indices, series_count, tokens_per_series = ( + dataset.token_mapping(dataset_index) + ) + np.savez_compressed( + attention_dir / f"{dataset_index:06d}_{safe_id}.npz", + attention=attention[start:end].float().cpu().numpy().astype(np.float16), + patch_logits=patch_logits[start:end].float().cpu().numpy().astype(np.float16), + original_token_indices=token_indices, + series_count=np.asarray(series_count, dtype=np.int32), + tokens_per_series=np.asarray(tokens_per_series, dtype=np.int32), + ) + + return np.concatenate(logits_out), np.concatenate(labels_out), indices_out + + +def make_scheduler( + optimizer: torch.optim.Optimizer, + total_steps: int, + warmup_fraction: float, +): + warmup_steps = int(total_steps * warmup_fraction) + + def factor(step: int) -> float: + if warmup_steps > 0 and step < warmup_steps: + return float(step + 1) / float(warmup_steps) + decay_steps = max(total_steps - warmup_steps, 1) + progress = min(max((step - warmup_steps) / decay_steps, 0.0), 1.0) + return 0.5 * (1.0 + math.cos(math.pi * progress)) + + return torch.optim.lr_scheduler.LambdaLR(optimizer, factor) + + +def main() -> None: + parser = argparse.ArgumentParser( + "NeuroVFM 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") + parser.add_argument("--epochs", type=int, default=50) + parser.add_argument("--batch_size", type=int, default=4) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument("--weight_decay", type=float, default=0.05) + parser.add_argument("--warmup_fraction", type=float, default=0.10) + parser.add_argument("--grad_clip", type=float, default=1.0) + parser.add_argument("--hidden_dim", type=int, default=512) + parser.add_argument("--mlp_hidden_dim", type=int, default=384) + parser.add_argument("--no_pos_weight", action="store_true") + parser.add_argument("--num_workers", type=int, default=0) + parser.add_argument("--save_test_attention", action="store_true", + help="Write per-study class attention and patch logits; this can be large.") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--device", type=str, + default="cuda" if torch.cuda.is_available() else "cpu", + ) + args = parser.parse_args() + + if args.epochs <= 0 or args.batch_size <= 0: + parser.error("--epochs and --batch_size must be positive") + if not 0.0 <= args.warmup_fraction <= 1.0: + parser.error("--warmup_fraction must be between 0 and 1") + if args.device.startswith("cuda") and not torch.cuda.is_available(): + parser.error("CUDA was requested but is not available") + + torch.manual_seed(args.seed) + np.random.seed(args.seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(args.seed) + + features_dir = Path(args.features_dir) + results_dir = Path(args.results_dir) + results_dir.mkdir(parents=True, exist_ok=True) + device = torch.device(args.device) + + label_names: list[str] = json.loads((features_dir / "label_names.json").read_text()) + if not label_names: + raise ValueError("label_names.json contains no classes") + train_data = RaggedTokenDataset(features_dir, "train") + val_data = RaggedTokenDataset(features_dir, "val") + test_data = RaggedTokenDataset(features_dir, "test") + fingerprints = { + dataset.metadata.get("cache_fingerprint") + for dataset in (train_data, val_data, test_data) + } + if None in fingerprints or len(fingerprints) != 1: + raise ValueError( + "Train/val/test token caches do not share one verified cache fingerprint" + ) + if not (train_data.dim == val_data.dim == test_data.dim): + raise ValueError("Feature dimensions differ across splits") + if not ( + train_data.labels.shape[1] + == val_data.labels.shape[1] + == test_data.labels.shape[1] + == len(label_names) + ): + raise ValueError("Label dimensions differ across splits or label_names.json") + split_ids = { + "train": set(train_data.subject_ids), + "val": set(val_data.subject_ids), + "test": set(test_data.subject_ids), + } + for left, right in (("train", "val"), ("train", "test"), ("val", "test")): + overlap = split_ids[left] & split_ids[right] + if overlap: + raise ValueError( + f"Study leakage between {left} and {right}: {len(overlap)} duplicate IDs" + ) + + pin_memory = device.type == "cuda" + train_loader = DataLoader( + train_data, batch_size=args.batch_size, shuffle=True, + num_workers=args.num_workers, pin_memory=pin_memory, + collate_fn=collate_ragged, drop_last=False, + ) + val_loader = DataLoader( + val_data, batch_size=args.batch_size, shuffle=False, + num_workers=args.num_workers, pin_memory=pin_memory, + collate_fn=collate_ragged, drop_last=False, + ) + test_loader = DataLoader( + test_data, batch_size=args.batch_size, shuffle=False, + num_workers=args.num_workers, pin_memory=pin_memory, + collate_fn=collate_ragged, drop_last=False, + ) + + n_classes = len(label_names) + head = ClassifyThenAggregate( + dim=train_data.dim, + n_classes=n_classes, + hidden_dim=args.hidden_dim, + mlp_hidden_dims=(args.mlp_hidden_dim,), + drop_rate=0.0, + use_gating=True, + use_norm=False, + use_output_bias_scale=True, + ).to(device) + print( + f"NeuroVFM CTA head: dim={train_data.dim}, attention_hidden={args.hidden_dim}, " + f"patch_mlp=[{args.mlp_hidden_dim}], classes={n_classes}" + ) + print( + f"Bags: train={len(train_data)}, val={len(val_data)}, test={len(test_data)} | " + f"tokens(train)={train_data.num_tokens:,}" + ) + + if args.no_pos_weight: + pos_weight = None + print("pos_weight: disabled") + else: + positives = train_data.labels.sum(axis=0) + negatives = len(train_data) - positives + weights = np.clip(negatives / np.clip(positives, 1.0, None), 1.0, 100.0) + pos_weight = torch.tensor(weights, dtype=torch.float32, device=device) + print(f"pos_weight: median={float(np.median(weights)):.1f}, max={float(weights.max()):.1f}") + + loss_fn = nn.BCEWithLogitsLoss(pos_weight=pos_weight) + optimizer = torch.optim.AdamW( + head.parameters(), lr=args.lr, weight_decay=args.weight_decay + ) + scheduler = make_scheduler( + optimizer, total_steps=max(len(train_loader) * args.epochs, 1), + warmup_fraction=args.warmup_fraction, + ) + + best_val_auroc = -1.0 + best_epoch = None + best_state: dict[str, torch.Tensor] | None = None + history: list[dict] = [] + + for epoch in range(1, args.epochs + 1): + head.train() + epoch_loss = 0.0 + n_seen = 0 + for tokens, labels, cu_seqlens, _indices in train_loader: + tokens = prepare_tokens(tokens, device) + labels = labels.to(device, non_blocking=True) + + optimizer.zero_grad(set_to_none=True) + with amp_context(device): + logits = head(tokens, cu_seqlens) + loss = loss_fn(logits.float(), labels) + loss.backward() + if args.grad_clip > 0: + nn.utils.clip_grad_norm_(head.parameters(), args.grad_clip) + optimizer.step() + scheduler.step() + + epoch_loss += loss.item() * labels.shape[0] + n_seen += labels.shape[0] + + train_loss = epoch_loss / max(n_seen, 1) + val_logits, val_labels, _ = evaluate_head(head, val_loader, device) + val_mean, _ = auroc_table(val_logits, val_labels, label_names) + history.append( + {"epoch": epoch, "train_loss": train_loss, "val_mean_auroc": val_mean} + ) + print( + f"epoch {epoch:3d} train_loss={train_loss:.4f} " + f"val_mean_AUROC={val_mean:.4f}" + ) + if np.isfinite(val_mean) and val_mean > best_val_auroc: + best_val_auroc = val_mean + best_epoch = epoch + best_state = { + name: value.detach().cpu().clone() + for name, value in head.state_dict().items() + } + + if best_state is None: + raise RuntimeError("No epoch produced a finite validation mean AUROC") + head.load_state_dict(best_state) + print(f"Best validation mean AUROC: {best_val_auroc:.4f}") + + best_val_logits, best_val_labels, _ = evaluate_head(head, val_loader, device) + validation_thresholds = select_validation_thresholds( + best_val_logits, best_val_labels + ) + + attention_dir = results_dir / "test_attention" if args.save_test_attention else None + test_logits, test_labels, test_indices = evaluate_head( + head, test_loader, device, attention_dir=attention_dir + ) + test_mean, test_per_class = auroc_table(test_logits, test_labels, label_names) + test_metric_rows = per_class_metrics( + test_logits, test_labels, label_names, validation_thresholds + ) + print(f"Test mean AUROC: {test_mean:.4f}") + + architecture = { + "name": "NeuroVFM ClassifyThenAggregate", + "dim": train_data.dim, + "n_classes": n_classes, + "hidden_dim": args.hidden_dim, + "mlp_hidden_dims": [args.mlp_hidden_dim], + "drop_rate": 0.0, + "init_std": 0.02, + "use_gating": True, + "use_norm": False, + "use_output_bias_scale": True, + } + torch.save( + { + "state_dict": best_state, + "architecture": architecture, + "label_names": label_names, + "cache_provenance": { + "train": train_data.metadata, + "val": val_data.metadata, + "test": test_data.metadata, + }, + "best_epoch": best_epoch, + "best_val_mean_auroc": best_val_auroc, + "validation_thresholds": validation_thresholds, + "args": vars(args), + }, + results_dir / "mil_head.pt", + ) + np.save(results_dir / "test_logits.npy", test_logits.astype(np.float32)) + np.save(results_dir / "test_labels.npy", test_labels.astype(np.float32)) + ordered_test_ids = [test_data.subject_ids[index] for index in test_indices] + (results_dir / "test_subject_ids.txt").write_text("\n".join(ordered_test_ids) + "\n") + (results_dir / "history.json").write_text(json.dumps(history, indent=2) + "\n") + (results_dir / "validation_thresholds.json").write_text( + json.dumps( + dict(zip(label_names, validation_thresholds.astype(float))), indent=2, + allow_nan=False, + ) + "\n" + ) + (results_dir / "per_class_test_auroc.json").write_text( + json.dumps( + { + "mean_auroc": test_mean, + "per_class": { + name: (value if np.isfinite(value) else None) + for name, value in test_per_class.items() + }, + "metrics": test_metric_rows, + }, + indent=2, + allow_nan=False, + ) + "\n" + ) + with (results_dir / "test_aurocs.csv").open("w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=list(test_metric_rows[0])) + writer.writeheader() + writer.writerows(test_metric_rows) + + print(f"Artifacts written to {results_dir}") + + +if __name__ == "__main__": + main() diff --git a/contrastive-pretraining/scripts/mil_probe_online.py b/contrastive-pretraining/scripts/mil_probe_online.py new file mode 100644 index 0000000..eba6906 --- /dev/null +++ b/contrastive-pretraining/scripts/mil_probe_online.py @@ -0,0 +1,961 @@ +#!/usr/bin/env python3 +"""Train the exact MIL head while encoding frozen MRI studies online. + +This entry point does not write or consume a token-embedding cache. The visual +encoder is reconstructed from the supplied checkpoint, frozen, and run under +torch.no_grad() for every study in every epoch. Only the shared +ClassifyThenAggregate head receives gradients. +""" + +from __future__ import annotations + +import argparse +import copy +import csv +import hashlib +import json +import math +import os +import random +import re +from contextlib import nullcontext +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch.utils.data import DataLoader + +from data_inference import MRReportDatasetInfer, collate_fn_infer +from extract_features import _load_and_verify, build_encoder +from mil_probe import ( + ClassifyThenAggregate, + auroc_table, + make_scheduler, + per_class_metrics, + select_validation_thresholds, +) + + +@dataclass +class DatasetMetadata: + subject_ids: list[str] + labels: np.ndarray + label_names: list[str] + + +@dataclass +class EncodedStudy: + tokens: torch.Tensor + target: torch.Tensor + subject_id: str + original_token_indices: np.ndarray | None + full_token_count: int + series_count: int + tokens_per_series: int + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Online frozen-encoder Classify-Then-Aggregate MIL" + ) + model = parser.add_argument_group("frozen encoder") + model.add_argument("--weights_path", required=True) + model.add_argument( + "--encoder", + default="vjepa2", + choices=("vjepa2", "vjepa21", "vjepa2_sliding", "vjepa21_sliding"), + ) + model.add_argument("--vjepa21_checkpoint", default=None) + model.add_argument("--chunk_size", type=int, default=64) + model.add_argument("--fusion_mode", default="late", choices=("late",)) + model.add_argument( + "--pooling_strategy", + default="simple_attn", + choices=("simple_attn", "cross_attn", "gated"), + ) + model.add_argument("--dim_latent", type=int, default=512) + model.add_argument("--extra_latent_projection", action="store_true") + model.add_argument("--strict_missing", action="store_true") + + data = parser.add_argument_group("study data") + data.add_argument("--data_folder", default=None) + data.add_argument("--jsonl_file", required=True) + data.add_argument("--labels_file", required=True) + data.add_argument("--splits_csv", required=True) + data.add_argument("--space", default="native_space") + data.add_argument( + "--normalizer", default="zscore", choices=("zscore", "percentile", "minmax") + ) + data.add_argument("--preprocessed_dir", default=None) + data.add_argument("--use_preprocessed", action="store_true") + data.add_argument("--cache_allow_mismatch", action="store_true") + data.add_argument( + "--max_tokens_per_study", + type=int, + default=0, + help="Deterministically subsample after encoding; 0 keeps every token.", + ) + + training = parser.add_argument_group("MIL training") + training.add_argument("--results_dir", default="./mil_probe_online_results") + training.add_argument("--epochs", type=int, default=50) + training.add_argument("--lr", type=float, default=5e-4) + training.add_argument("--weight_decay", type=float, default=0.05) + training.add_argument("--warmup_fraction", type=float, default=0.10) + training.add_argument("--grad_clip", type=float, default=1.0) + training.add_argument("--grad_accum_steps", type=int, default=4) + training.add_argument("--hidden_dim", type=int, default=512) + training.add_argument("--mlp_hidden_dim", type=int, default=384) + training.add_argument("--no_pos_weight", action="store_true") + training.add_argument("--num_workers", type=int, default=4) + training.add_argument("--seed", type=int, default=42) + training.add_argument( + "--amp_dtype", + default="auto", + choices=("auto", "bfloat16", "float16", "float32"), + ) + training.add_argument("--resume", default=None, help="Path to an online last.pt checkpoint.") + training.add_argument("--save_test_attention", action="store_true") + return parser.parse_args() + + +def validate_args(args: argparse.Namespace) -> None: + if not torch.cuda.is_available(): + raise RuntimeError("Online MIL currently requires CUDA because encoder construction is CUDA-only.") + if args.epochs <= 0 or args.lr <= 0: + raise ValueError("--epochs and --lr must be positive") + if args.weight_decay < 0 or args.grad_clip < 0: + raise ValueError("--weight_decay and --grad_clip cannot be negative") + if not 0.0 <= args.warmup_fraction <= 1.0: + raise ValueError("--warmup_fraction must be in [0, 1]") + if args.grad_accum_steps <= 0: + raise ValueError("--grad_accum_steps must be positive") + if args.hidden_dim <= 0 or args.mlp_hidden_dim <= 0: + raise ValueError("MIL hidden dimensions must be positive") + if args.num_workers < 0 or args.max_tokens_per_study < 0: + raise ValueError("Worker and token counts cannot be negative") + if args.chunk_size <= 0 or args.dim_latent <= 0: + raise ValueError("--chunk_size and --dim_latent must be positive") + if not Path(args.weights_path).is_file(): + raise FileNotFoundError(f"Encoder checkpoint not found: {args.weights_path}") + for name in ("jsonl_file", "labels_file", "splits_csv"): + value = Path(getattr(args, name)) + if not value.is_file(): + raise FileNotFoundError(f"--{name} not found: {value}") + if args.use_preprocessed: + if not args.preprocessed_dir or not Path(args.preprocessed_dir).is_dir(): + raise ValueError("--use_preprocessed requires an existing --preprocessed_dir") + elif not args.data_folder or not Path(args.data_folder).is_dir(): + raise ValueError("Raw online encoding requires an existing --data_folder") + if args.vjepa21_checkpoint and not Path(args.vjepa21_checkpoint).is_file(): + raise FileNotFoundError(f"--vjepa21_checkpoint not found: {args.vjepa21_checkpoint}") + if args.resume and not Path(args.resume).is_file(): + raise FileNotFoundError(f"Resume checkpoint not found: {args.resume}") + + +def seed_everything(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def seed_worker(worker_id: int) -> None: + del worker_id + worker_seed = torch.initial_seed() % (2**32) + random.seed(worker_seed) + np.random.seed(worker_seed) + + +def resolve_compute_dtype(name: str) -> torch.dtype: + if name == "auto": + supports_bf16 = getattr(torch.cuda, "is_bf16_supported", lambda: False)() + return torch.bfloat16 if supports_bf16 else torch.float16 + if name == "bfloat16": + supports_bf16 = getattr(torch.cuda, "is_bf16_supported", lambda: False)() + if not supports_bf16: + raise RuntimeError("This CUDA device does not support bfloat16; use --amp_dtype float16.") + return torch.bfloat16 + if name == "float16": + return torch.float16 + return torch.float32 + + +def autocast_context(dtype: torch.dtype): + if dtype == torch.float32: + return nullcontext() + return torch.autocast(device_type="cuda", dtype=dtype) + + +def make_grad_scaler(enabled: bool): + try: + return torch.amp.GradScaler("cuda", enabled=enabled) + except (AttributeError, TypeError): + return torch.cuda.amp.GradScaler(enabled=enabled) + + +def sha256_file(path: str | Path) -> str: + digest = hashlib.sha256() + with Path(path).open("rb") as handle: + for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def file_record(path: str | Path | None) -> dict[str, Any] | None: + if path is None: + return None + resolved = Path(path).resolve() + record: dict[str, Any] = {"path": str(resolved), "exists": resolved.exists()} + if resolved.is_file(): + stat = resolved.stat() + record.update({"size": stat.st_size, "sha256": sha256_file(resolved)}) + return record + + +def cohort_digest(metadata: DatasetMetadata) -> str: + digest = hashlib.sha256() + digest.update(json.dumps(metadata.subject_ids, separators=(",", ":")).encode("utf-8")) + digest.update(np.ascontiguousarray(metadata.labels, dtype=np.float32).tobytes()) + digest.update(json.dumps(metadata.label_names, separators=(",", ":")).encode("utf-8")) + return digest.hexdigest() + + +def fingerprint(value: Any) -> str: + payload = json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def build_dataset(args: argparse.Namespace, split: str) -> MRReportDatasetInfer: + return MRReportDatasetInfer( + data_folder=args.data_folder, + jsonl_file=args.jsonl_file, + space=args.space, + normalizer=args.normalizer, + labels_file=args.labels_file, + splits_csv=args.splits_csv, + split=split, + preprocessed_dir=args.preprocessed_dir, + use_preprocessed=args.use_preprocessed, + cache_allow_mismatch=args.cache_allow_mismatch, + ) + + +def dataset_metadata(dataset: MRReportDatasetInfer, split: str) -> DatasetMetadata: + if len(dataset) == 0: + raise ValueError(f"The {split} split contains no eligible studies") + label_names = [str(name) for name in dataset.label_columns] + if not label_names: + raise ValueError("The labels file does not define any output classes") + subject_ids: list[str] = [] + vectors: list[np.ndarray] = [] + missing: list[str] = [] + for sample in dataset.samples: + subject_id = str(sample["subject_id"]) + if "labels" not in sample: + missing.append(subject_id) + continue + vector = np.asarray(sample["labels"], dtype=np.float32) + if vector.shape != (len(label_names),): + raise ValueError( + f"{split} study {subject_id!r} has label shape {vector.shape}; " + f"expected {(len(label_names),)}" + ) + subject_ids.append(subject_id) + vectors.append(vector) + if missing: + preview = ", ".join(missing[:5]) + raise ValueError(f"The {split} split has {len(missing)} unlabeled studies: {preview}") + if len(subject_ids) != len(dataset.samples): + raise RuntimeError(f"Internal metadata alignment failed for the {split} split") + if len(subject_ids) != len(set(subject_ids)): + raise ValueError(f"Duplicate subject IDs detected within the {split} split") + labels = np.stack(vectors, axis=0).astype(np.float32, copy=False) + if not np.isfinite(labels).all(): + raise ValueError(f"The {split} split contains non-finite labels") + if not np.all((labels == 0.0) | (labels == 1.0)): + raise ValueError(f"The {split} split contains labels outside strict binary values 0 and 1") + return DatasetMetadata(subject_ids=subject_ids, labels=labels, label_names=label_names) + + +def validate_split_pair( + left: DatasetMetadata, + left_name: str, + right: DatasetMetadata, + right_name: str, +) -> None: + if left.label_names != right.label_names: + raise ValueError(f"Label schema mismatch between {left_name} and {right_name}") + overlap = sorted(set(left.subject_ids).intersection(right.subject_ids)) + if overlap: + preview = ", ".join(overlap[:5]) + raise ValueError( + f"Subject leakage between {left_name} and {right_name}: " + f"{len(overlap)} overlapping IDs ({preview})" + ) + + +def make_loader( + dataset: MRReportDatasetInfer, + *, + shuffle: bool, + num_workers: int, + generator: torch.Generator | None = None, +) -> DataLoader: + return DataLoader( + dataset, + batch_size=1, + shuffle=shuffle, + num_workers=num_workers, + drop_last=False, + collate_fn=collate_fn_infer, + pin_memory=True, + persistent_workers=False, + worker_init_fn=seed_worker, + generator=generator, + ) + + +def evenly_spaced_indices(length: int, maximum: int) -> np.ndarray: + if maximum <= 0 or length <= maximum: + return np.arange(length, dtype=np.int64) + return np.rint(np.linspace(0, length - 1, maximum)).astype(np.int64) + + +def encode_study( + encoder: torch.nn.Module, + batch: tuple[Any, ...], + device: torch.device, + compute_dtype: torch.dtype, + max_tokens_per_study: int, + *, + keep_mapping: bool, +) -> EncodedStudy: + images, _sentences, subject_id, real_volume_mask, labels = batch + if not isinstance(images, torch.Tensor) or images.ndim != 6 or images.shape[0] != 1: + raise ValueError("Online MIL expects images with shape [1,R,1,D,H,W]") + images = images.to(device=device, dtype=compute_dtype, non_blocking=True) + real_volume_mask = torch.as_tensor(real_volume_mask, dtype=torch.bool, device=device) + if real_volume_mask.ndim != 2 or real_volume_mask.shape[0] != 1: + raise ValueError("real_volume_mask must have shape [1,R]") + with torch.no_grad(): + with autocast_context(compute_dtype): + visual_tokens, token_mask = encoder( + text_input=None, + image=images, + device=device, + real_volume_mask=real_volume_mask, + return_loss=False, + return_visual_tokens=True, + ) + if visual_tokens.ndim != 3 or token_mask.shape != visual_tokens.shape[:2]: + raise RuntimeError("Encoder output must be [B,N,D] plus a matching [B,N] mask") + valid_tokens = visual_tokens[0, token_mask[0].bool()].detach() + full_token_count = int(valid_tokens.shape[0]) + if full_token_count <= 0: + raise ValueError(f"Study {subject_id!r} produced no valid visual tokens") + series_count = int(real_volume_mask[0].sum().item()) + if series_count <= 0 or full_token_count % series_count != 0: + raise RuntimeError( + f"Cannot map {full_token_count} valid tokens across {series_count} valid series" + ) + tokens_per_series = full_token_count // series_count + selected_indices = evenly_spaced_indices(full_token_count, max_tokens_per_study) + if selected_indices.size != full_token_count: + index_tensor = torch.from_numpy(selected_indices).to(device=valid_tokens.device) + valid_tokens = valid_tokens.index_select(0, index_tensor) + if not torch.isfinite(valid_tokens).all(): + raise ValueError(f"Study {subject_id!r} produced non-finite visual tokens") + if valid_tokens.requires_grad: + raise RuntimeError("Frozen encoder output unexpectedly requires gradients") + target = torch.as_tensor(labels, dtype=torch.float32, device=device).reshape(1, -1) + return EncodedStudy( + tokens=valid_tokens, + target=target, + subject_id=str(subject_id), + original_token_indices=selected_indices if keep_mapping else None, + full_token_count=full_token_count, + series_count=series_count, + tokens_per_series=tokens_per_series, + ) + + +def cumulative_lengths(token_count: int) -> torch.Tensor: + return torch.tensor([0, token_count], dtype=torch.long) + + +def verify_gradient_boundary(encoder: torch.nn.Module, head: torch.nn.Module) -> None: + if any(parameter.grad is not None for parameter in encoder.parameters()): + raise RuntimeError("Frozen encoder accumulated gradients") + if not any(parameter.grad is not None for parameter in head.parameters()): + raise RuntimeError("MIL head received no gradients") + + +def train_one_epoch( + encoder: torch.nn.Module, + head: ClassifyThenAggregate, + loader: DataLoader, + criterion: torch.nn.Module, + optimizer: torch.optim.Optimizer, + scheduler: Any, + scaler: Any, + device: torch.device, + compute_dtype: torch.dtype, + grad_accum_steps: int, + grad_clip: float, + max_tokens_per_study: int, + global_step: int, + gradient_boundary_checked: bool, +) -> tuple[float, int, bool]: + encoder.eval() + head.train() + optimizer.zero_grad(set_to_none=True) + total_loss = 0.0 + total_studies = len(loader) + for batch_index, batch in enumerate(loader): + study = encode_study( + encoder, batch, device, compute_dtype, max_tokens_per_study, keep_mapping=False + ) + if study.target.shape[1] != head.n_classes: + raise ValueError( + f"Study {study.subject_id!r} has {study.target.shape[1]} labels; " + f"the MIL head expects {head.n_classes}" + ) + block_start = (batch_index // grad_accum_steps) * grad_accum_steps + block_size = min(grad_accum_steps, total_studies - block_start) + with autocast_context(compute_dtype): + logits = head(study.tokens, cumulative_lengths(study.tokens.shape[0])) + loss = criterion(logits.float(), study.target) + if not torch.isfinite(loss): + raise FloatingPointError(f"Non-finite training loss for study {study.subject_id!r}") + total_loss += float(loss.detach().cpu()) + scaler.scale(loss / block_size).backward() + if not gradient_boundary_checked: + verify_gradient_boundary(encoder, head) + gradient_boundary_checked = True + at_boundary = (batch_index + 1) % grad_accum_steps == 0 + at_end = batch_index + 1 == total_studies + if at_boundary or at_end: + if scaler.is_enabled(): + scaler.unscale_(optimizer) + if grad_clip > 0: + torch.nn.utils.clip_grad_norm_(head.parameters(), grad_clip) + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad(set_to_none=True) + scheduler.step() + global_step += 1 + return total_loss / total_studies, global_step, gradient_boundary_checked + + +def safe_subject_id(subject_id: str) -> str: + cleaned = re.sub(r"[^A-Za-z0-9_.-]+", "_", subject_id).strip("._") + return cleaned or "study" + + +def atomic_savez(path: Path, **arrays: Any) -> None: + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp.npz") + np.savez_compressed(temporary, **arrays) + os.replace(temporary, path) + + +def evaluate_online( + encoder: torch.nn.Module, + head: ClassifyThenAggregate, + loader: DataLoader, + device: torch.device, + compute_dtype: torch.dtype, + max_tokens_per_study: int, + attention_dir: Path | None = None, +) -> tuple[np.ndarray, np.ndarray, list[str]]: + encoder.eval() + head.eval() + logits_out: list[np.ndarray] = [] + labels_out: list[np.ndarray] = [] + subject_ids: list[str] = [] + if attention_dir is not None: + if attention_dir.exists() and any(attention_dir.iterdir()): + raise FileExistsError(f"Refusing to mix attention exports in {attention_dir}") + attention_dir.mkdir(parents=True, exist_ok=True) + with torch.no_grad(): + for dataset_index, batch in enumerate(loader): + study = encode_study( + encoder, + batch, + device, + compute_dtype, + max_tokens_per_study, + keep_mapping=attention_dir is not None, + ) + cu_seqlens = cumulative_lengths(study.tokens.shape[0]) + with autocast_context(compute_dtype): + if attention_dir is None: + logits = head(study.tokens, cu_seqlens) + else: + logits, attention, patch_logits = head( + study.tokens, cu_seqlens, return_details=True + ) + logits_out.append(logits.float().cpu().numpy()) + labels_out.append(study.target.float().cpu().numpy()) + subject_ids.append(study.subject_id) + if attention_dir is not None: + if study.original_token_indices is None: + raise RuntimeError("Attention export is missing token-index metadata") + output_path = attention_dir / ( + f"{dataset_index:06d}_{safe_subject_id(study.subject_id)}.npz" + ) + atomic_savez( + output_path, + attention=attention.float().cpu().numpy().astype(np.float16), + patch_logits=patch_logits.float().cpu().numpy().astype(np.float16), + original_token_indices=study.original_token_indices, + full_token_count=np.asarray(study.full_token_count, dtype=np.int64), + series_count=np.asarray(study.series_count, dtype=np.int32), + tokens_per_series=np.asarray(study.tokens_per_series, dtype=np.int32), + ) + if not logits_out: + raise ValueError("Cannot evaluate an empty study loader") + return np.concatenate(logits_out), np.concatenate(labels_out), subject_ids + + +def cpu_state_dict(module: torch.nn.Module) -> dict[str, torch.Tensor]: + return {name: value.detach().cpu().clone() for name, value in module.state_dict().items()} + + +def rng_state(train_generator: torch.Generator) -> dict[str, Any]: + return { + "python": random.getstate(), + "numpy": np.random.get_state(), + "torch": torch.get_rng_state(), + "cuda": torch.cuda.get_rng_state_all(), + "train_generator": train_generator.get_state(), + } + + +def restore_rng_state(state: dict[str, Any], train_generator: torch.Generator) -> None: + random.setstate(state["python"]) + np.random.set_state(state["numpy"]) + torch.set_rng_state(state["torch"]) + torch.cuda.set_rng_state_all(state["cuda"]) + train_generator.set_state(state["train_generator"]) + + +def atomic_torch_save(value: Any, path: Path) -> None: + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + torch.save(value, temporary) + os.replace(temporary, path) + + +def atomic_save_npy(path: Path, array: np.ndarray) -> None: + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + with temporary.open("wb") as handle: + np.save(handle, array) + os.replace(temporary, path) + + +def atomic_write_text(path: Path, text: str) -> None: + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + temporary.write_text(text, encoding="utf-8") + os.replace(temporary, path) + + +def load_torch_checkpoint(path: str | Path, device: torch.device) -> dict[str, Any]: + try: + return torch.load(path, map_location=device, weights_only=False) + except TypeError: + return torch.load(path, map_location=device) + + +def json_safe(value: Any) -> Any: + if isinstance(value, dict): + return {str(key): json_safe(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [json_safe(item) for item in value] + if isinstance(value, np.ndarray): + return json_safe(value.tolist()) + if isinstance(value, (np.floating, float)): + numeric = float(value) + return numeric if math.isfinite(numeric) else None + if isinstance(value, np.integer): + return int(value) + return value + + +def make_provenance( + args: argparse.Namespace, + train_metadata: DatasetMetadata, + val_metadata: DatasetMetadata, +) -> dict[str, Any]: + return { + "schema_version": 1, + "mode": "online_frozen_encoder_mil", + "encoder_checkpoint": file_record(args.weights_path), + "vjepa21_checkpoint": file_record(args.vjepa21_checkpoint), + "encoder": { + "name": args.encoder, + "chunk_size": args.chunk_size, + "fusion_mode": args.fusion_mode, + "pooling_strategy": args.pooling_strategy, + "dim_latent": args.dim_latent, + "extra_latent_projection": args.extra_latent_projection, + "strict_missing": args.strict_missing, + }, + "data": { + "data_folder": str(Path(args.data_folder).resolve()) if args.data_folder else None, + "jsonl_file": file_record(args.jsonl_file), + "labels_file": file_record(args.labels_file), + "splits_csv": file_record(args.splits_csv), + "space": args.space, + "normalizer": args.normalizer, + "use_preprocessed": args.use_preprocessed, + "preprocessed_dir": ( + str(Path(args.preprocessed_dir).resolve()) if args.preprocessed_dir else None + ), + "cache_allow_mismatch": args.cache_allow_mismatch, + }, + "cohorts": { + "train": {"count": len(train_metadata.subject_ids), "digest": cohort_digest(train_metadata)}, + "val": {"count": len(val_metadata.subject_ids), "digest": cohort_digest(val_metadata)}, + }, + "label_names": train_metadata.label_names, + } + + +def immutable_description( + args: argparse.Namespace, + provenance: dict[str, Any], + dim_latent: int, + n_classes: int, + optimizer_steps_per_epoch: int, + resolved_dtype: torch.dtype, +) -> dict[str, Any]: + return { + "provenance": provenance, + "architecture": { + "dim": dim_latent, + "n_classes": n_classes, + "hidden_dim": args.hidden_dim, + "mlp_hidden_dim": args.mlp_hidden_dim, + "drop_rate": 0.0, + "use_gating": True, + "use_norm": False, + "use_output_bias_scale": True, + }, + "training": { + "epochs": args.epochs, + "lr": args.lr, + "weight_decay": args.weight_decay, + "warmup_fraction": args.warmup_fraction, + "grad_clip": args.grad_clip, + "grad_accum_steps": args.grad_accum_steps, + "optimizer_steps_per_epoch": optimizer_steps_per_epoch, + "no_pos_weight": args.no_pos_weight, + "seed": args.seed, + "amp_dtype": args.amp_dtype, + "resolved_dtype": str(resolved_dtype), + "max_tokens_per_study": args.max_tokens_per_study, + }, + } + + +def save_csv(path: Path, rows: list[dict[str, Any]]) -> None: + if not rows: + raise ValueError("Cannot write an empty metrics CSV") + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + with temporary.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + os.replace(temporary, path) + + +def main() -> None: + args = parse_args() + validate_args(args) + seed_everything(args.seed) + results_dir = Path(args.results_dir) + if results_dir.exists() and any(results_dir.iterdir()) and args.resume is None: + raise FileExistsError( + f"Refusing to overwrite nonempty results directory: {results_dir}. " + "Use a new directory or --resume." + ) + results_dir.mkdir(parents=True, exist_ok=True) + device = torch.device("cuda") + compute_dtype = resolve_compute_dtype(args.amp_dtype) + + train_dataset = build_dataset(args, "train") + val_dataset = build_dataset(args, "val") + train_metadata = dataset_metadata(train_dataset, "train") + val_metadata = dataset_metadata(val_dataset, "val") + validate_split_pair(train_metadata, "train", val_metadata, "val") + train_generator = torch.Generator().manual_seed(args.seed) + train_loader = make_loader( + train_dataset, shuffle=True, num_workers=args.num_workers, generator=train_generator + ) + val_loader = make_loader(val_dataset, shuffle=False, num_workers=args.num_workers) + + print("Building frozen visual encoder") + encoder, dim_latent = build_encoder(args) + _load_and_verify(encoder, args.weights_path, strict_missing=args.strict_missing) + try: + image_encoder = encoder.visual_transformer + if hasattr(image_encoder, "model") and hasattr(image_encoder.model, "merge_and_unload"): + image_encoder.model.merge_and_unload() + except Exception as error: + print(f"LoRA merge skipped: {error}") + encoder.requires_grad_(False) + encoder.to(device=device, dtype=compute_dtype) + encoder.eval() + if any(parameter.requires_grad for parameter in encoder.parameters()): + raise RuntimeError("Failed to freeze every visual-encoder parameter") + + n_classes = len(train_metadata.label_names) + head = ClassifyThenAggregate( + dim=dim_latent, + n_classes=n_classes, + hidden_dim=args.hidden_dim, + mlp_hidden_dims=(args.mlp_hidden_dim,), + drop_rate=0.0, + init_std=0.02, + use_gating=True, + use_norm=False, + use_output_bias_scale=True, + ).to(device) + if args.no_pos_weight: + pos_weight = None + else: + positives = train_metadata.labels.sum(axis=0) + negatives = len(train_metadata.labels) - positives + weights = np.clip(negatives / np.maximum(positives, 1.0), 1.0, 100.0) + pos_weight = torch.as_tensor(weights, dtype=torch.float32, device=device) + criterion = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight) + optimizer = torch.optim.AdamW(head.parameters(), lr=args.lr, weight_decay=args.weight_decay) + optimizer_steps_per_epoch = math.ceil(len(train_loader) / args.grad_accum_steps) + scheduler = make_scheduler( + optimizer, optimizer_steps_per_epoch * args.epochs, args.warmup_fraction + ) + scaler = make_grad_scaler(compute_dtype == torch.float16) + provenance = make_provenance(args, train_metadata, val_metadata) + immutable = immutable_description( + args, + provenance, + dim_latent, + n_classes, + optimizer_steps_per_epoch, + compute_dtype, + ) + run_fingerprint = fingerprint(immutable) + + start_epoch = 0 + global_step = 0 + best_epoch = -1 + best_val_mean_auroc = float("-inf") + best_state: dict[str, torch.Tensor] | None = None + history: list[dict[str, Any]] = [] + gradient_boundary_checked = False + if args.resume is not None: + resume = load_torch_checkpoint(args.resume, device) + if resume.get("run_fingerprint") != run_fingerprint: + raise ValueError( + "Resume checkpoint does not match the encoder, data, head, or training configuration" + ) + head.load_state_dict(resume["head_state_dict"]) + optimizer.load_state_dict(resume["optimizer_state_dict"]) + scheduler.load_state_dict(resume["scheduler_state_dict"]) + scaler.load_state_dict(resume["scaler_state_dict"]) + start_epoch = int(resume["completed_epoch"]) + 1 + global_step = int(resume["global_step"]) + best_epoch = int(resume["best_epoch"]) + best_val_mean_auroc = float(resume["best_val_mean_auroc"]) + best_state = resume.get("best_state_dict") + history = list(resume["history"]) + gradient_boundary_checked = bool(resume.get("gradient_boundary_checked", False)) + restore_rng_state(resume["rng_state"], train_generator) + if start_epoch > args.epochs: + raise ValueError("Resume checkpoint is beyond the configured number of epochs") + print(f"Resuming at epoch {start_epoch + 1} of {args.epochs}") + + for epoch in range(start_epoch, args.epochs): + train_loss, global_step, gradient_boundary_checked = train_one_epoch( + encoder, + head, + train_loader, + criterion, + optimizer, + scheduler, + scaler, + device, + compute_dtype, + args.grad_accum_steps, + args.grad_clip, + args.max_tokens_per_study, + global_step, + gradient_boundary_checked, + ) + val_logits, val_labels, _ = evaluate_online( + encoder, + head, + val_loader, + device, + compute_dtype, + args.max_tokens_per_study, + ) + val_mean_auroc, val_per_class = auroc_table( + val_logits, val_labels, train_metadata.label_names + ) + history.append( + json_safe( + { + "epoch": epoch + 1, + "train_loss": train_loss, + "val_mean_auroc": val_mean_auroc, + "val_per_class_auroc": val_per_class, + "optimizer_steps": global_step, + "lr": optimizer.param_groups[0]["lr"], + } + ) + ) + if math.isfinite(val_mean_auroc) and val_mean_auroc > best_val_mean_auroc: + best_val_mean_auroc = float(val_mean_auroc) + best_epoch = epoch + 1 + best_state = cpu_state_dict(head) + atomic_torch_save( + { + "state_dict": best_state, + "best_epoch": best_epoch, + "best_val_mean_auroc": best_val_mean_auroc, + "label_names": train_metadata.label_names, + "run_fingerprint": run_fingerprint, + "provenance": provenance, + }, + results_dir / "best.pt", + ) + atomic_torch_save( + { + "head_state_dict": head.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + "scheduler_state_dict": scheduler.state_dict(), + "scaler_state_dict": scaler.state_dict(), + "completed_epoch": epoch, + "global_step": global_step, + "best_epoch": best_epoch, + "best_val_mean_auroc": best_val_mean_auroc, + "best_state_dict": best_state, + "history": history, + "gradient_boundary_checked": gradient_boundary_checked, + "rng_state": rng_state(train_generator), + "run_fingerprint": run_fingerprint, + "immutable_run": immutable, + "args": vars(args), + }, + results_dir / "last.pt", + ) + print( + f"Epoch {epoch + 1:03d}/{args.epochs:03d} " + f"train_loss={train_loss:.6f} val_mean_auroc={val_mean_auroc:.6f}" + ) + + if best_state is None or best_epoch < 1: + raise RuntimeError( + "Validation AUROC was non-finite in every epoch; no test evaluation was performed" + ) + head.load_state_dict(best_state) + head.to(device) + best_val_logits, best_val_labels, _ = evaluate_online( + encoder, + head, + val_loader, + device, + compute_dtype, + args.max_tokens_per_study, + ) + validation_thresholds = select_validation_thresholds(best_val_logits, best_val_labels) + if not np.isfinite(validation_thresholds).all(): + raise RuntimeError("Validation threshold selection produced non-finite values") + + test_dataset = build_dataset(args, "test") + test_metadata = dataset_metadata(test_dataset, "test") + validate_split_pair(train_metadata, "train", test_metadata, "test") + validate_split_pair(val_metadata, "val", test_metadata, "test") + test_loader = make_loader(test_dataset, shuffle=False, num_workers=args.num_workers) + attention_dir = results_dir / "test_attention" if args.save_test_attention else None + test_logits, test_labels, test_subject_ids = evaluate_online( + encoder, + head, + test_loader, + device, + compute_dtype, + args.max_tokens_per_study, + attention_dir=attention_dir, + ) + if test_subject_ids != test_metadata.subject_ids: + raise RuntimeError("Test DataLoader order does not match validated test metadata") + test_mean_auroc, test_per_class = auroc_table( + test_logits, test_labels, train_metadata.label_names + ) + test_metric_rows = per_class_metrics( + test_logits, test_labels, train_metadata.label_names, validation_thresholds + ) + + final_provenance = copy.deepcopy(provenance) + final_provenance["cohorts"]["test"] = { + "count": len(test_metadata.subject_ids), + "digest": cohort_digest(test_metadata), + } + architecture = { + "name": "ClassifyThenAggregate", + "dim": dim_latent, + "n_classes": n_classes, + "hidden_dim": args.hidden_dim, + "mlp_hidden_dims": [args.mlp_hidden_dim], + "drop_rate": 0.0, + "init_std": 0.02, + "use_gating": True, + "use_norm": False, + "use_output_bias_scale": True, + } + atomic_torch_save( + { + "state_dict": cpu_state_dict(head), + "architecture": architecture, + "label_names": train_metadata.label_names, + "data_provenance": final_provenance, + "best_epoch": best_epoch, + "best_val_mean_auroc": best_val_mean_auroc, + "validation_thresholds": validation_thresholds.tolist(), + "args": vars(args), + }, + results_dir / "mil_head.pt", + ) + atomic_save_npy(results_dir / "test_logits.npy", test_logits.astype(np.float32)) + atomic_save_npy(results_dir / "test_labels.npy", test_labels.astype(np.float32)) + atomic_write_text(results_dir / "test_subject_ids.txt", "\n".join(test_subject_ids) + "\n") + atomic_write_text( + results_dir / "history.json", + json.dumps(json_safe(history), indent=2, allow_nan=False) + "\n", + ) + threshold_payload = dict(zip(train_metadata.label_names, validation_thresholds.tolist())) + atomic_write_text( + results_dir / "validation_thresholds.json", + json.dumps(json_safe(threshold_payload), indent=2, allow_nan=False) + "\n", + ) + metric_payload = { + "mean_auroc": test_mean_auroc, + "per_class": test_per_class, + "metrics": test_metric_rows, + } + atomic_write_text( + results_dir / "per_class_test_auroc.json", + json.dumps(json_safe(metric_payload), indent=2, allow_nan=False) + "\n", + ) + save_csv(results_dir / "test_aurocs.csv", test_metric_rows) + print(f"Best validation epoch: {best_epoch}") + print(f"Best validation mean AUROC: {best_val_mean_auroc:.6f}") + print(f"Test mean AUROC: {test_mean_auroc:.6f}") + print(f"Results written to {results_dir.resolve()}") + + +if __name__ == "__main__": + main()