diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py index 2b4a0dc30..e24a9d036 100644 --- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py @@ -10,10 +10,6 @@ MortalityPredictionStageNetMIMIC4: ICD codes + 10-dim lab vectors, patient-level samples aggregated across all admissions. ---task icd_labs - ICDLabsMIMIC4: ICD codes + 10-dim lab vectors via the unified - multimodal pipeline. No notes required. - --task clinical_notes_icd_labs ClinicalNotesICDLabsMIMIC4: discharge/radiology notes + ICD + labs. Requires --note-root. Legacy; ICD codes are discharge-coded (leakage). @@ -54,7 +50,9 @@ from __future__ import annotations import argparse +import inspect import csv +import warnings from pathlib import Path from typing import Any, Tuple @@ -71,17 +69,21 @@ split_by_sample, ) from pyhealth.models import MLP, RNN, Transformer, UnifiedMultimodalEmbeddingModel +from pyhealth.models.embedding import VisionEmbeddingModel from pyhealth.models.bottleneck_transformer import BottleneckTransformer from pyhealth.models.ehrmamba import EHRMamba from pyhealth.models.jamba_ehr import JambaEHR +from pyhealth.processors import fit_lab_standardizer, lab_standardizer_fit_scope from pyhealth.tasks import MortalityPredictionStageNetMIMIC4 from pyhealth.tasks.multimodal_mimic4 import ( ClinicalNotesICDLabsMIMIC4, ICDLabsMIMIC4, + LabsOnlyMIMIC4, NotesLabsMIMIC4, + CXRMultimodalMIMIC4, ) from pyhealth.trainer import Trainer -from pyhealth.utils import set_seed +from pyhealth.utils import set_seed, write_run_config def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset: @@ -96,11 +98,45 @@ def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset: if args.task == "icd_labs": ehr_tables = ["diagnoses_icd", "procedures_icd", "labevents"] + if args.task in ("notes_labs", "notes_only"): + if not args.note_root: + raise ValueError(f"--task {args.task} requires --note-root.") + note_tables = [getattr(args, "note_source", "discharge")] + # Load ICD tables only when explicitly requested (they are discharge-coded). + ehr_tables = ( + ["diagnoses_icd", "procedures_icd", "labevents"] + if args.icd_codes + else ["labevents"] + ) + if args.include_vitals: + if "chartevents" not in ehr_tables: + ehr_tables.append("chartevents") + + if args.task == "labs_only": + # Pure EHR baseline: only labevents, no notes, no ICD codes. + ehr_tables = ["labevents"] + note_tables = None + + cxr_tables = None + if args.task in ("cxr_only", "cxr_labs", "cxr_notes_labs"): + if not args.cxr_root: + raise ValueError(f"--task {args.task} requires --cxr-root.") + # ``metadata`` supplies image_path, StudyDate/StudyTime, and ViewPosition. + cxr_tables = ["metadata"] + ehr_tables = ["labevents"] if args.task != "cxr_only" else [] + if args.task == "cxr_notes_labs": + if not args.note_root: + raise ValueError("--task cxr_notes_labs requires --note-root.") + note_tables = [getattr(args, "note_source", "discharge")] + return MIMIC4Dataset( ehr_root=args.ehr_root, ehr_tables=ehr_tables, note_root=args.note_root if note_tables else None, note_tables=note_tables, + cxr_root=args.cxr_root if cxr_tables else None, + cxr_tables=cxr_tables, + cxr_variant=args.cxr_variant, cache_dir=args.cache_dir, dev=args.dev if args.dev else False, num_workers=args.num_workers, @@ -114,27 +150,149 @@ def _build_task(args: argparse.Namespace): return ICDLabsMIMIC4(window_hours=args.observation_window_hours) if args.task == "clinical_notes_icd_labs": return ClinicalNotesICDLabsMIMIC4(window_hours=args.observation_window_hours) - if args.task == "notes_labs": - return NotesLabsMIMIC4( + if args.task in ("notes_labs", "notes_only"): + # Pass only what this checkout's task class accepts. A flag the class + # cannot honour must stop the run, not be dropped: a silently ignored + # --discharge-note-policy would record one protocol in run_config.json + # and execute another. + wanted = dict( window_hours=args.observation_window_hours, include_icd=args.icd_codes, include_vitals=args.include_vitals, + include_labs=(args.task != "notes_only"), + note_extraction=getattr(args, "note_extraction", "regex"), + note_source=getattr(args, "note_source", "discharge"), + discharge_note_policy=getattr( + args, "discharge_note_policy", "extraction"), + text_normalize=getattr(args, "text_normalize", "none"), + ) + accepted = set( + inspect.signature(NotesLabsMIMIC4.__init__).parameters + ) + defaults = { + "include_icd": False, + "include_vitals": False, + "include_labs": True, + "note_extraction": "regex", + "note_source": "discharge", + "discharge_note_policy": "extraction", + "text_normalize": "none", + } + unsupported = [ + name + for name, value in wanted.items() + if name not in accepted and value != defaults.get(name) + ] + if unsupported: + raise SystemExit( + f"NotesLabsMIMIC4 in this checkout does not accept " + f"{', '.join(sorted(unsupported))}. Either drop the flag or use " + f"a checkout whose task class supports it." + ) + task_kwargs = {k: v for k, v in wanted.items() if k in accepted} + if args.task == "notes_only" and "include_labs" not in accepted: + raise SystemExit( + "--task notes_only needs a NotesLabsMIMIC4 that accepts " + "include_labs; this checkout always emits labs." + ) + task = NotesLabsMIMIC4(**task_kwargs) + if args.tokenizer_model: + schema_key = "admission_note_times" + _, opts = task.input_schema[schema_key] + task.input_schema[schema_key] = ("tuple_time_text", {**opts, "tokenizer_model": args.tokenizer_model}) + print(f"[tokenizer] Overriding tokenizer_model → {args.tokenizer_model}") + # Note token budget. The processor default (128) truncates ~96% of + # extracted discharge notes, so the encoder sees only their first ~20%. + # input_schema is part of the task cache key, so each budget caches apart. + _nml = getattr(args, "note_max_length", None) + if _nml: + # Fail loudly rather than let HF silently clamp to the encoder's + # position-embedding limit (BERT/Bio_ClinicalBERT = 512). + _tokmodel = args.tokenizer_model or task.input_schema[ + "admission_note_times"][1].get("tokenizer_model", "") + if _nml > 512 and "longformer" not in _tokmodel.lower(): + raise SystemExit( + f"--note-max-length {_nml} exceeds the 512 position-embedding " + f"limit of {_tokmodel!r}. Use --tokenizer-model " + f"yikuan8/Clinical-Longformer for budgets >512." + ) + schema_key = "admission_note_times" + _, opts = task.input_schema[schema_key] + task.input_schema[schema_key] = ( + "tuple_time_text", {**opts, "max_length": _nml} + ) + print(f"[tokenizer] note max_length → {_nml}") + return task + if args.task == "labs_only": + return LabsOnlyMIMIC4(window_hours=args.observation_window_hours) + if args.task in ("cxr_only", "cxr_labs", "cxr_notes_labs"): + return CXRMultimodalMIMIC4( + window_hours=args.observation_window_hours, + include_labs=args.task in ("cxr_labs", "cxr_notes_labs"), + include_notes=args.task == "cxr_notes_labs", + frontal_only=not args.cxr_all_views, + image_size=args.cxr_image_size, + max_images=args.cxr_max_images, + note_source=args.note_source, ) raise ValueError(f"Unknown task: {args.task}") -def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]: +def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any, str]: + """Split by patient, falling back to by-sample only if that yields nothing. + + The fallback is leaky: a patient with several admissions can then land in + both train and test, which inflates the metrics. It only triggers on tiny + cohorts, but it must not trigger silently, so the mode is returned and + recorded alongside the run's results. + """ train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed) if len(train_ds) == 0 or len(test_ds) == 0: + warnings.warn( + "split_by_patient produced an empty split, falling back to " + "split_by_sample. The same patient may now appear in train and " + "test, so these metrics are optimistic and not comparable to " + "patient-split runs.", + RuntimeWarning, + stacklevel=2, + ) train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed) - return train_ds, val_ds, test_ds - - -def _build_model(args: argparse.Namespace, sample_dataset: Any): + return train_ds, val_ds, test_ds, "by_sample_fallback_leaky" + return train_ds, val_ds, test_ds, "by_patient" + + +def _resolve_finetune_mode(args: argparse.Namespace) -> str: + """--freeze-encoder is a back-compat alias for --text-finetune-mode frozen.""" + return "frozen" if args.freeze_encoder else args.text_finetune_mode + + +def _build_model( + args: argparse.Namespace, + sample_dataset: Any, + numeric_standardizers: dict[str, Any] | None = None, +): + finetune_mode = _resolve_finetune_mode(args) + field_embeddings = None + if "cxr" in sample_dataset.input_processors: + # Reuse the repository's VisionEmbeddingModel in the unified image + # branch; only the temporal alignment is supplied by unified.py. + field_embeddings = { + "cxr": VisionEmbeddingModel( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + patch_size=16, + backbone="patch", + pretrained=False, + ) + } unified = UnifiedMultimodalEmbeddingModel( processors=sample_dataset.input_processors, embedding_dim=args.embedding_dim, - freeze_text_encoder=args.freeze_encoder, + field_embeddings=field_embeddings, + freeze_text_encoder=(finetune_mode == "frozen"), + normalize_content=not getattr(args, "no_normalize_content", False), + numeric_standardizers=numeric_standardizers, + cache_frozen_text=not getattr(args, "no_text_cache", False), ) if args.model == "mlp": @@ -228,6 +386,108 @@ def _write_predictions( ) +def _load_pretrained_weights(model, ckpt_path: str) -> None: + """Load SSL pretraining weights into a supervised model. + + Delegates to ``BaseModel.load_pretrained_state_dict``, which performs the + architecture-specific key mapping (downstream models name the unified + backbone ``_unified_backbone`` / ``_unified_jamba`` / ``_unified_blocks``) + and REQUIRES full backbone coverage. + + The previous implementation here built ``model.state_dict()``, overwrote + whichever checkpoint keys happened to map, and called ``strict=False``. + Because the untouched tensors were already present, PyTorch reported no + missing keys, so a jamba/mamba checkpoint that matched only ~6 of ~30 + backbone tensors trained on a largely RANDOM backbone while looking + perfectly healthy. That silently corrupted real Table 2 cells. + """ + print(f"[pretrain] Loading checkpoint from {ckpt_path}") + state = torch.load(ckpt_path, map_location="cpu", weights_only=True) + try: + stats = model.load_pretrained_state_dict(state) + except ValueError as exc: + # Surface this as a run-level failure naming the checkpoint. A partial + # unified backbone must abort the job rather than train on random + # weights, and the operator needs to know WHICH checkpoint was bad. + raise RuntimeError( + f"Refusing to train on a partial unified backbone from {ckpt_path}: {exc}" + ) from exc + print( + "[pretrain] backbone {}/{} tensors, embedding {} matched; " + "uninitialised: {}".format( + stats["backbone_matched"], stats["backbone_target"], + stats["embedding_matched"], stats["missing_keys"][:6], + ) + ) + +def _note_availability_report(train_ds, sample_limit: int = 4000) -> dict: + """Measure how often a note is actually present, and whether that leaks. + + A missing note is represented by a fixed placeholder embedding, so + "has a real note" is trivially learnable. Measured on MIMIC-IV, mortality was + 5.67% where a note existed against 1.37% where it did not, a 4.1x gap: note + AVAILABILITY carries outcome signal with no clinical content behind it. That + confound has to be visible on every run rather than rediscovered, so it is + measured on the TRAIN split (never test) and recorded in run_config.json. + """ + import numpy as np + + total = len(train_ds) + if total == 0: + return {} + # Stride across the split rather than taking a prefix: samples are grouped by + # patient, so the first N are not representative of note availability. + step = max(1, total // sample_limit) + indices = range(0, total, step) + present, labels = [], [] + for i in indices: + try: + sample = train_ds[i] + except Exception: + break + field = sample.get("admission_note_times") + if field is None: + return {} + try: + mask = field["mask"] if isinstance(field, dict) else field[1] + mask = torch.as_tensor(mask) + if mask.ndim == 1: + mask = mask.unsqueeze(0) + # content tokens = attention mask less [CLS] and [SEP] + content = int((mask.sum(dim=1) - 2).clamp(min=0).max()) + except Exception: + return {} + present.append(content > 5) + labels.append(float(sample.get("mortality", 0))) + if not present or all(present) or not any(present): + return {"note_present_rate": float(np.mean(present)) if present else None, + "n_inspected": len(present)} + present = np.array(present); labels = np.array(labels) + report = { + "n_inspected": int(len(present)), + "note_present_rate": float(present.mean()), + "mortality_with_note": float(labels[present].mean()), + "mortality_without_note": float(labels[~present].mean()), + } + ratio = (report["mortality_with_note"] / + max(report["mortality_without_note"], 1e-9)) + report["mortality_ratio_present_vs_absent"] = float(ratio) + print( + f"[note-availability] {100*report['note_present_rate']:.1f}% of train " + f"samples carry a real note; mortality {report['mortality_with_note']:.4f} " + f"with vs {report['mortality_without_note']:.4f} without ({ratio:.1f}x)." + ) + if ratio > 1.5 or ratio < 0.67: + warnings.warn( + f"Note availability is {ratio:.1f}x associated with the outcome. The " + "missing-note placeholder is a constant embedding, so this is " + "learnable signal with no clinical content. Report it, restrict to " + "complete cases, or model missingness explicitly.", + RuntimeWarning, stacklevel=2, + ) + return report + + def _compute_pos_weight(train_ds, label_key: str = "mortality") -> float: """Count pos/neg in train_ds and return n_neg/n_pos for BCE pos_weight.""" n_pos = n_neg = 0 @@ -260,7 +520,56 @@ def run(args: argparse.Namespace) -> Path: "Task produced zero samples. Check roots/tables or adjust settings." ) - train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) + split_seed = getattr(args, "split_seed", None) + split_seed_pinned = split_seed is not None + if split_seed is None: + split_seed = args.seed + # Measured at full scale on labs_only: letting the split follow the seed + # gives sd(PR-AUC) 0.0236, versus 0.0042 with the split pinned, so test-set + # composition contributes ~31x the variance of initialisation. Cells + # compared across different splits are unpaired and far less sensitive. + warnings.warn( + "--split-seed not set, so the patient split follows --seed and every " + "seed draws a different test set. For an ablation, pin --split-seed " + "across compared cells: it is roughly a 5x sensitivity gain for no " + "extra compute.", + RuntimeWarning, stacklevel=2, + ) + print(f"[split] patient split seed={split_seed} " + f"({'pinned' if split_seed_pinned else 'follows --seed'}); " + f"training seed={args.seed}") + train_ds, val_ds, test_ds, split_mode = _split_dataset(sample_dataset, seed=split_seed) + if getattr(args, "pretrained_ckpt", None) and split_mode != "by_patient": + raise ValueError( + "Pretrained comparisons require a non-empty patient-level split. " + "Refusing the sample-level fallback because the checkpoint's train " + "statistics could then include a downstream test patient." + ) + + # Fit before any outcome-dependent resampling and exclusively on the train + # subset. ``SampleDataset.subset`` exposes only its selected indices here; + # LabStandardizer iterates this object, never ``sample_dataset``. + note_availability = _note_availability_report(train_ds) + + numeric_standardizers: dict[str, Any] = {} + if "labs" in sample_dataset.input_processors and not getattr( + args, "no_lab_standardization", False + ): + if "labs_mask" not in sample_dataset.input_processors: + raise RuntimeError("Lab standardisation requires the labs_mask observation field.") + lab_standardizer = fit_lab_standardizer( + train_ds, + value_field="labs", + fit_scope=lab_standardizer_fit_scope(train_ds, value_field="labs"), + ) + numeric_standardizers["labs"] = lab_standardizer + print( + "[lab-standardization] fitted on train split only: " + f"counts={lab_standardizer.observed_count.tolist()} " + f"mean={lab_standardizer.mean.tolist()} std={lab_standardizer.std.tolist()}" + ) + elif "labs" in sample_dataset.input_processors: + print("[lab-standardization] disabled; reproducing raw-lab baseline.") label_key = list(sample_dataset.output_schema.keys())[0] @@ -272,22 +581,40 @@ def run(args: argparse.Namespace) -> Path: if strategy == "undersample": ratio = args.balanced_ratio - print(f"[sampling] Undersampling negatives -> pos:neg 1:{ratio}") + print(f"[sampling] Undersampling negatives → pos:neg 1:{ratio}") train_ds = sample_balanced(train_ds, ratio=ratio, seed=args.seed, label_key=label_key) print(f"[sampling] Training size after undersample: {len(train_ds)}") elif strategy == "oversample": ratio = args.balanced_ratio - print(f"[sampling] Oversampling positives -> pos:neg 1:{ratio}") + print(f"[sampling] Oversampling positives → pos:neg 1:{ratio}") train_ds = sample_oversample(train_ds, ratio=ratio, seed=args.seed, label_key=label_key) print(f"[sampling] Training size after oversample: {len(train_ds)}") elif strategy == "weighted": - print("[sampling] Weighted resampling (class-proportional, with replacement)") + print("[sampling] Weighted resampling (class-proportional, with replacement, no external sampler)") train_ds = sample_weighted(train_ds, seed=args.seed, label_key=label_key) print(f"[sampling] Training size after weighted resample: {len(train_ds)}") - model = _build_model(args, sample_dataset) + model = _build_model(args, sample_dataset, numeric_standardizers) + + # Load pretrained SSL weights if requested. + if getattr(args, "pretrained_ckpt", None): + _load_pretrained_weights(model, args.pretrained_ckpt) + if args.numeric_input_stats_path: + model.embedding_model.capture_numeric_encoder_input_stats( + args.numeric_input_stats_path, field_name="labs" + ) + # Use the same deterministic train-split batch as the pretraining + # audit. This records the true projection input while avoiding a + # misleading difference caused only by independent shuffling. + audit_batch = next(iter(get_dataloader( + train_ds, batch_size=args.batch_size, shuffle=False, + ))) + with torch.no_grad(): + model(**audit_batch) + if not Path(args.numeric_input_stats_path).is_file(): + raise RuntimeError("Numeric-input audit hook did not produce its artifact.") # Apply class-imbalance correction via BCE pos_weight. # pos_weight = n_neg / n_pos so the rare positive class gets proportionally @@ -300,25 +627,103 @@ def run(args: argparse.Namespace) -> Path: print(f"[pos_weight] Using pos_weight={pw_value:.2f} for binary BCE loss.") model._pos_weight = torch.tensor([pw_value], dtype=torch.float32) - train_loader = get_dataloader(train_ds, batch_size=args.batch_size, shuffle=True) + # DataLoader worker options arrive with the performance PR. Pass only what + # this checkout's get_dataloader accepts, and stop if the caller asked for + # one it cannot honour, so a requested option is never silently ignored. + _wanted_loader = { + "num_workers": args.loader_num_workers, + "pin_memory": args.pin_memory, + "persistent_workers": ( + args.persistent_workers and args.loader_num_workers > 0 + ), + "prefetch_factor": ( + args.prefetch_factor if args.loader_num_workers > 0 else None + ), + } + _loader_accepts = set(inspect.signature(get_dataloader).parameters) + _loader_defaults = { + "num_workers": 0, + "pin_memory": False, + "persistent_workers": False, + "prefetch_factor": None, + } + _unsupported_loader = [ + name + for name, value in _wanted_loader.items() + if name not in _loader_accepts and value != _loader_defaults[name] + ] + if _unsupported_loader: + raise SystemExit( + f"get_dataloader in this checkout does not accept " + f"{', '.join(sorted(_unsupported_loader))}. Drop the flag, or use a " + f"checkout that includes the DataLoader worker options." + ) + loader_kwargs = { + k: v for k, v in _wanted_loader.items() if k in _loader_accepts + } + train_loader = get_dataloader( + train_ds, + batch_size=args.batch_size, + shuffle=True, + **loader_kwargs, + ) val_loader = ( - get_dataloader(val_ds, batch_size=args.batch_size, shuffle=False) + get_dataloader( + val_ds, + batch_size=args.batch_size, + shuffle=False, + **loader_kwargs, + ) if len(val_ds) > 0 else None ) test_loader = ( - get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) + get_dataloader( + test_ds, + batch_size=args.batch_size, + shuffle=False, + **loader_kwargs, + ) if len(test_ds) > 0 else None ) - # Experiment name encodes model + seed for easy log separation - exp_name = f"{args.model}_seed{args.seed}" + # Which split the reported predictions come from. Falling back to val, or + # worse to train, silently reports held-in performance as if it were test, + # so resolve it once here and record it alongside the results. + if test_loader is not None: + inference_loader, eval_split = test_loader, "test" + elif val_loader is not None: + inference_loader, eval_split = val_loader, "val" + warnings.warn( + "No test split available; reporting predictions from the VALIDATION " + "split. These are not test metrics.", + RuntimeWarning, stacklevel=2, + ) + else: + inference_loader, eval_split = train_loader, "train" + warnings.warn( + "No test or validation split available; reporting predictions from " + "the TRAINING split. These metrics are held-in and meaningless as a " + "generalisation estimate.", + RuntimeWarning, stacklevel=2, + ) + + # The task MUST be in the name. Without it, two arms of the same comparison + # at the same seed, for example --task labs_only and --task notes_labs, + # resolve to one directory and the second run overwrites the first run's + # metrics_history.json, run_config.json and predictions CSV. The loss is + # silent: the surviving directory looks like a complete run. + exp_name = f"{args.task}_{args.model}_seed{args.seed}" output_dir = Path(args.output_dir) trainer = Trainer( model=model, - metrics=["pr_auc", "roc_auc", "f1", "f1_opt", "accuracy"], + # f1_opt is a threshold-optimised F1 that this checkout's + # binary_metrics_fn does not implement, and requesting it aborts + # validation at the end of epoch 1. Model selection uses pr_auc, which + # is a rank metric and needs no threshold. + metrics=["pr_auc", "roc_auc", "f1", "accuracy"], device=args.device, enable_logging=True, output_path=str(output_dir), @@ -352,7 +757,37 @@ def run(args: argparse.Namespace) -> Path: optimizer_params["lr"] = effective_lr + # Record the resolved conditions, not the raw flags: text_finetune_mode and + # the learning rate are both derived, so the CLI alone does not identify the + # run. Without this the artifacts cannot say whether the encoder was frozen. + write_run_config( + str(output_dir / exp_name), + { + **vars(args), + "resolved_text_finetune_mode": _resolve_finetune_mode(args), + "resolved_lr": effective_lr, + "resolved_max_grad_norm": effective_max_grad_norm, + "split_mode": split_mode, + "eval_split": eval_split, + "split_seed_pinned": split_seed_pinned, + "note_availability": note_availability, + "n_train": len(train_ds), + "n_val": len(val_ds), + "n_test": len(test_ds), + }, + ) + if args.epochs > 0 and len(train_ds) > 0: + # PR-AUC/ROC-AUC are undefined for a single-class validation fold. + # Full MIMIC patient splits contain both labels; tiny real-data demo + # runs do not, so use finite validation loss for checkpoint selection. + val_labels = { + int(float(val_ds[i][label_key])) for i in range(len(val_ds)) + } + monitor = "pr_auc" if len(val_labels) == 2 else "loss" + monitor_criterion = "max" if monitor == "pr_auc" else "min" + if monitor != "pr_auc": + print("[monitor] Validation fold has one class; selecting checkpoints by loss.") trainer.train( train_dataloader=train_loader, val_dataloader=val_loader, @@ -360,12 +795,17 @@ def run(args: argparse.Namespace) -> Path: optimizer_params=optimizer_params, weight_decay=args.weight_decay, max_grad_norm=effective_max_grad_norm, - monitor="pr_auc", + monitor=monitor, + monitor_criterion=monitor_criterion, load_best_model_at_last=True, patience=args.patience, + encoder_lr=args.encoder_lr, + # Mixed precision is a property of the training loop, not of the + # Trainer object. + use_amp=args.use_amp, + amp_dtype=args.amp_dtype, ) - inference_loader = test_loader or val_loader or train_loader y_true, y_prob, _, patient_ids = trainer.inference( inference_loader, return_patient_ids=True ) @@ -381,13 +821,41 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument("--ehr-root", type=str, required=True) parser.add_argument("--note-root", type=str, default=None) + parser.add_argument("--cxr-root", type=str, default=None) + parser.add_argument( + "--cxr-variant", + choices=["default", "sunlab"], + default="default", + help=( + "Layout of the CXR root. 'default' reads " + "mimic-cxr-2.0.0-metadata-pyhealth.csv and needs a " + "studytime_normalized column. 'sunlab' reads the resized set, " + "normalises StudyTime itself and derives image paths from " + "dicom_id. Choosing the wrong one fails at dataset build with " + "KeyError: 'studytime_normalized'." + ), + ) parser.add_argument("--cache-dir", type=str, default=None) parser.add_argument("--output-dir", type=str, default="./output/unified_e2e") + parser.add_argument( + "--numeric-input-stats-path", type=str, default=None, + help="Optional JSON artifact: first lab tensor entering the numeric encoder.", + ) + parser.add_argument( + "--pretrained-ckpt", + type=str, + default=None, + help=( + "Path to a SSL pretraining checkpoint (e.g., from " + "scripts/pretrain_ssl.py). Loads embedding_model and backbone " + "weights into the downstream model." + ), + ) parser.add_argument( "--task", type=str, - choices=["icd_labs", "clinical_notes_icd_labs"], + choices=["stagenet", "icd_labs", "clinical_notes_icd_labs", "notes_labs", "notes_only", "labs_only", "cxr_only", "cxr_labs", "cxr_notes_labs"], default="stagenet", help=( "notes_labs: admission-context text (CC/HPI/PMH/MedsOnAdm) + labs. " @@ -430,7 +898,32 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--weight-decay", type=float, default=0.0) parser.add_argument("--device", type=str, default=None) parser.add_argument("--num-workers", type=int, default=1) + parser.add_argument( + "--loader-num-workers", + type=int, + default=0, + help="Worker processes for runtime batch loading/collation.", + ) + parser.add_argument("--pin-memory", action="store_true", default=False) + parser.add_argument("--persistent-workers", action="store_true", default=False) + parser.add_argument("--prefetch-factor", type=int, default=2) + amp_group = parser.add_mutually_exclusive_group() + amp_group.add_argument("--use-amp", dest="use_amp", action="store_true") + amp_group.add_argument("--no-amp", dest="use_amp", action="store_false") + parser.set_defaults(use_amp=False) + parser.add_argument( + "--amp-dtype", + choices=["bf16", "fp16"], + default="bf16", + help="AMP compute dtype; bf16 is recommended on A100/H100.", + ) parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--split-seed", + type=int, + default=None, + help="Patient split seed. Defaults to --seed for backward compatibility.", + ) parser.add_argument("--patience", type=int, default=None) parser.add_argument( "--dev", @@ -457,6 +950,12 @@ def parse_args() -> argparse.Namespace: # Task-specific parser.add_argument("--observation-window-hours", type=int, default=24) + parser.add_argument( + "--cxr-all-views", action="store_true", default=False, + help="Keep non-frontal CXR views; default restricts to PA/AP frontal images.", + ) + parser.add_argument("--cxr-image-size", type=int, default=224) + parser.add_argument("--cxr-max-images", type=int, default=4) parser.add_argument( "--icd-codes", action="store_true", @@ -468,6 +967,22 @@ def parse_args() -> argparse.Namespace: "Enable only for ablation / legacy comparison experiments." ), ) + parser.add_argument( + "--discharge-note-policy", + choices=["extraction", "charttime"], + default="extraction", + help="How the observation window applies to discharge summaries. " + "extraction (default, Lee et al. 2023): retrieve across the " + "admission and let admission-context section extraction be the " + "temporal control. charttime: strictly non-anticipative, but " + "drops the summary for ~90%% of admissions.", + ) + parser.add_argument( + "--no-text-cache", + action="store_true", + help="Disable the frozen-text [CLS] cache. Diagnostic: isolates the cache " + "as a cause when a run fails to optimise.", + ) parser.add_argument( "--freeze-encoder", action="store_true", @@ -475,8 +990,44 @@ def parse_args() -> argparse.Namespace: help=( "Freeze pretrained BERT text encoder weights and train only the " "downstream backbone (MLP/RNN/Transformer head + projection layer). " - "Reduces VRAM by ~50% for the text branch; useful when GPU memory " - "is limited or for faster iteration on backbone architectures." + "Reduces VRAM by ~50%% for the text branch; useful when GPU memory " + "is limited or for faster iteration on backbone architectures. " + "Back-compat alias for --text-finetune-mode frozen." + ), + ) + parser.add_argument( + "--text-finetune-mode", + type=str, + default="full", + help=( + "Text encoder fine-tuning regime: 'full' (train all encoder params), " + "'frozen' (train only the head/projection), 'topk:N' (unfreeze the top " + "N transformer layers, embeddings stay frozen), or 'lora:r' (rank-r " + "LoRA adapters on attention, base frozen; needs peft). " + "Overridden by --freeze-encoder when that flag is set." + ), + ) + parser.add_argument( + "--encoder-lr", + type=float, + default=None, + help=( + "Discriminative learning rate for the pretrained text encoder. When " + "set, the encoder trains at this LR while the projection + downstream " + "head keep the base --lr. Recommended ~2e-5 for full/topk fine-tuning " + "to avoid destabilizing the encoder. Default None = single global LR." + ), + ) + parser.add_argument( + "--note-source", + type=str, + default="discharge", + choices=["discharge", "radiology"], + help=( + "Which MIMIC note table to use for notes_labs. 'discharge' (default) " + "uses admission-context discharge sections; 'radiology' uses radiology " + "report Impression/Findings, concatenated per admission (mirrors the " + "multimodal-EHR benchmark). Pair with a RadBERT --tokenizer-model." ), ) parser.add_argument( @@ -508,6 +1059,83 @@ def parse_args() -> argparse.Namespace: "Default: 1.0 (equal pos/neg). Used with undersample and oversample strategies." ), ) + parser.add_argument( + "--tokenizer-model", + type=str, + default=None, + help=( + "Override the tokenizer/encoder for notes. Must be a HuggingFace model ID. " + "Default: None (uses task class default, emilyalsentzer/Bio_ClinicalBERT). " + "Changes the task cache UUID so different tokenizers use isolated caches. " + "Examples: microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext, " + "medicalai/ClinicalBERT, yikuan8/Clinical-Longformer." + ), + ) + parser.add_argument( + "--no-normalize-content", + action="store_true", + default=False, + help=( + "Disable content normalization in the unified embedding, reproducing " + "pre-repair behaviour where text events were ~94%% patient-independent " + "constant while raw labs dominated. Use only to reproduce old runs." + ), + ) + parser.add_argument( + "--no-lab-standardization", + action="store_true", + default=False, + help=( + "Disable train-split-only per-analyte lab z-scoring. Use only for " + "the raw-lab ablation; default standardises immediately before the " + "numeric projection and checkpoints the fitted statistics." + ), + ) + parser.add_argument( + "--note-max-length", + type=int, + default=None, + help=( + "Note tokenization budget (TupleTimeTextProcessor.max_length; default " + "128, which truncates ~96%% of extracted discharge notes). BERT-family " + "encoders cap at 512 position embeddings; for longer budgets pass " + "--tokenizer-model yikuan8/Clinical-Longformer (4096) and expect to " + "need --batch-size 1." + ), + ) + parser.add_argument( + "--text-normalize", + type=str, + default="none", + choices=["none", "punct", "stopwords", "both"], + help=( + "Cut tokens per note before tokenization. 'punct': strip punctuation " + "(decimal points and thousands separators inside numbers are kept, so " + "lab values survive). 'stopwords': drop common English stopwords. " + "'both': both. Changes task_name so each setting gets its own cache." + ), + ) + parser.add_argument( + "--note-extraction", + type=str, + default="regex", + choices=[ + "regex", "regex_priority", "compact", "tfidf", + "section_hpi", "section_cc", "section_pmh", "section_meds", + "section_social", "section_family", "section_allergies", "section_ros", + "lab_retrieval", + ], + help=( + "Note text extraction strategy. " + "'regex' (default): section headers, document order. " + "'regex_priority': HPI-first order (fixes 48%% truncation rate). " + "'compact': HPI + CC only — always fits in 512 tokens. " + "'tfidf': TF overlap paragraph retrieval, no headers required. " + "'section_': single-section ablation (hpi/cc/pmh/meds/" + "social/family/allergies/ros). " + "Non-regex values isolate the task cache UUID." + ), + ) parser.add_argument( "--sampling-strategy", type=str, @@ -518,7 +1146,7 @@ def parse_args() -> argparse.Namespace: "'none': no resampling (default). " "'undersample': drop majority-class (neg) samples via sample_balanced(). " "'oversample': duplicate minority-class (pos) samples via sample_oversample(). " - "'weighted': class-proportional resampling w/ replacement via sample_weighted(). " + "'weighted': WeightedRandomSampler for batch-level balance without dataset modification. " "--balanced-sampling is a legacy alias for 'undersample'." ), ) @@ -528,7 +1156,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--rnn-layers", type=int, default=1) parser.add_argument("--bidirectional", action="store_true") - # Transformer / BottleneckTransformer shared + # Transformer / BottleneckTransformer shared. + # Standardized compute: 64-dim, 1 layer, 2 heads (head_dim 32) across all archs. parser.add_argument("--heads", type=int, default=4) parser.add_argument("--num-layers", type=int, default=2) @@ -553,9 +1182,11 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--mamba-conv-kernel", type=int, default=4, help="Causal conv kernel size for EHRMamba and JambaEHR blocks.") parser.add_argument("--jamba-transformer-layers", type=int, default=2, - help="Number of Transformer (attention) layers in JambaEHR.") - parser.add_argument("--jamba-mamba-layers", type=int, default=6, - help="Number of Mamba (SSM) layers in JambaEHR.") + help="Number of Transformer (attention) layers in JambaEHR. " + "Standard: 1 (a single Jamba block = 1 attn + 1 mamba).") + parser.add_argument("--jamba-mamba-layers", type=int, default=2, + help="Number of Mamba (SSM) layers in JambaEHR. " + "Default 2, matching the library JambaEHR default.") return parser.parse_args() diff --git a/pyhealth/datasets/configs/mimic4_cxr.yaml b/pyhealth/datasets/configs/mimic4_cxr.yaml index fb5045eef..bca6eb95e 100644 --- a/pyhealth/datasets/configs/mimic4_cxr.yaml +++ b/pyhealth/datasets/configs/mimic4_cxr.yaml @@ -5,7 +5,7 @@ tables: patient_id: "subject_id" timestamp: - "studydate" - - "studytime" + - "studytime_normalized" timestamp_format: "%Y%m%d%H%M%S" attributes: - "image_path" diff --git a/pyhealth/datasets/mimic4.py b/pyhealth/datasets/mimic4.py index 089a6d2f3..db47697b4 100644 --- a/pyhealth/datasets/mimic4.py +++ b/pyhealth/datasets/mimic4.py @@ -191,6 +191,15 @@ def __init__( log_memory_usage(f"After initializing {dataset_name}") def prepare_metadata(self, root: str) -> None: + prepared_path = os.path.join(root, "mimic-cxr-2.0.0-metadata-pyhealth.csv") + # The prepared file holds absolute JPEG paths. To rewrite its 100+ MB + # CSV for every CXR experiment is needless shared-filesystem IO, and it + # races a concurrent reader, so reuse the file when it is present. An + # installation with raw metadata only still uses the builder below. + if os.path.exists(prepared_path): + header = pd.read_csv(prepared_path, nrows=1) + if "image_path" in header.columns: + return metadata = pd.read_csv( os.path.join(root, "mimic-cxr-2.0.0-metadata.csv.gz"), dtype=str ) @@ -217,9 +226,7 @@ def process_image_path(x): metadata["image_path"] = metadata.apply(process_image_path, axis=1) - metadata.to_csv( - os.path.join(root, "mimic-cxr-2.0.0-metadata-pyhealth.csv"), index=False - ) + metadata.to_csv(prepared_path, index=False) return @@ -228,7 +235,7 @@ class MIMIC4CXRSunlabDataset(BaseDataset): Sunlab variant of the MIMIC-CXR Chest X-ray dataset. This variant uses the existing metadata CSV and derives flattened image - paths at ``images/{dicom_id}.jpg``. + paths at ``{images|resized_images}/{dicom_id}.jpg``. """ def __init__( @@ -245,7 +252,11 @@ def __init__( os.path.dirname(__file__), "configs", "mimic4_cxr_sunlab.yaml" ) logger.info(f"Using default Sunlab CXR config: {config_path}") - self.prepare_metadata(root) + metadata_csv = self.prepare_metadata(root, cache_dir=cache_dir) + if os.path.dirname(os.path.abspath(metadata_csv)) != os.path.abspath(root): + config_path = self._rewrite_sunlab_config( + config_path, metadata_csv, cache_dir or os.path.dirname(metadata_csv) + ) log_memory_usage(f"Before initializing {dataset_name}") super().__init__( root=root, @@ -267,7 +278,26 @@ def _resolve_column_name(columns: List[str], target: str) -> str: ) return resolved - def prepare_metadata(self, root: str) -> None: + @staticmethod + def _rewrite_sunlab_config( + config_path: str, metadata_csv: str, dest_dir: str + ) -> str: + """Point the sunlab YAML at a metadata CSV that is not under root.""" + with open(config_path, encoding="utf-8") as f: + text = f.read() + rewritten = text.replace( + "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv", + metadata_csv, + ) + os.makedirs(dest_dir, exist_ok=True) + out = os.path.join(dest_dir, "mimic4_cxr_sunlab.generated.yaml") + with open(out, "w", encoding="utf-8") as f: + f.write(rewritten) + return out + + def prepare_metadata( + self, root: str, cache_dir: Optional[str] = None + ) -> str: metadata_path = os.path.join(root, "mimic-cxr-2.0.0-metadata.csv") if not os.path.exists(metadata_path): raise FileNotFoundError( @@ -275,12 +305,25 @@ def prepare_metadata(self, root: str) -> None: "Expected existing metadata linked by dicom_id/subject_id/study_id." ) - images_dir = os.path.join(root, "images") - if not os.path.isdir(images_dir): + # The flattened layout appears under more than one directory name + # depending on how the set was produced, so accept either rather than + # hardcoding one and failing on a complete, correct dataset. + candidates = ("images", "resized_images") + images_dir = next( + ( + os.path.join(root, name) + for name in candidates + if os.path.isdir(os.path.join(root, name)) + ), + None, + ) + if images_dir is None: raise FileNotFoundError( - f"Sunlab images directory not found: {images_dir}. " - "Expected flattened image files at images/{dicom_id}.jpg." + f"No flattened image directory under {root}. Looked for " + f"{', '.join(candidates)}, each expected to hold " + "{dicom_id}.jpg." ) + images_subdir = os.path.basename(images_dir) metadata = pd.read_csv(metadata_path, dtype=str) @@ -307,15 +350,35 @@ def normalize_studytime(value: Optional[str]) -> str: metadata[study_time_col] = metadata[study_time_col].apply(normalize_studytime) metadata["image_path"] = metadata[dicom_col].apply( - lambda dicom_id: os.path.join(root, "images", f"{dicom_id}.jpg") + lambda dicom_id: os.path.join(root, images_subdir, f"{dicom_id}.jpg") ) # Align with existing config conventions by using lowercase headers. metadata.columns = [col.lower() for col in metadata.columns] - metadata.to_csv( - os.path.join(root, "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv"), - index=False, + filename = "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv" + dest_dirs = [] + if cache_dir: + dest_dirs.append(str(cache_dir)) + dest_dirs.append(root) + + for d in dest_dirs: + existing = os.path.join(d, filename) + if os.path.isfile(existing): + return existing + + last_err: Optional[OSError] = None + for d in dest_dirs: + os.makedirs(d, exist_ok=True) + dest = os.path.join(d, filename) + try: + metadata.to_csv(dest, index=False) + return dest + except OSError as exc: + last_err = exc + continue + raise PermissionError( + f"Could not write {filename} under {dest_dirs}: {last_err}" ) diff --git a/pyhealth/models/jamba_ehr.py b/pyhealth/models/jamba_ehr.py index 37d738f3f..7b8062eae 100644 --- a/pyhealth/models/jamba_ehr.py +++ b/pyhealth/models/jamba_ehr.py @@ -73,7 +73,7 @@ class JambaLayer(nn.Module): Args: feature_size (int): Hidden dimension shared by all layers. num_transformer_layers (int): Number of attention layers. Default 2. - num_mamba_layers (int): Number of SSM layers. Default 6. + num_mamba_layers (int): Number of SSM layers. Default 2. heads (int): Attention heads for Transformer layers. Default 4. dropout (float): Dropout rate for Transformer layers. Default 0.3. state_size (int): SSM state size for Mamba layers. Default 16. @@ -94,7 +94,7 @@ def __init__( self, feature_size: int, num_transformer_layers: int = 2, - num_mamba_layers: int = 6, + num_mamba_layers: int = 2, heads: int = 4, dropout: float = 0.3, state_size: int = 16, @@ -187,7 +187,7 @@ class JambaEHR(BaseModel): dataset (SampleDataset): Dataset providing processed inputs. embedding_dim (int): Embedding and hidden dimension. Default 128. num_transformer_layers (int): Transformer layers per stream. Default 2. - num_mamba_layers (int): Mamba layers per stream. Default 6. + num_mamba_layers (int): Mamba layers per stream. Default 2. heads (int): Attention heads per Transformer block. Default 4. dropout (float): Dropout rate. Default 0.3. state_size (int): SSM state size in Mamba blocks. Default 16. @@ -237,7 +237,7 @@ def __init__( dataset: SampleDataset, embedding_dim: int = 128, num_transformer_layers: int = 2, - num_mamba_layers: int = 6, + num_mamba_layers: int = 2, heads: int = 4, dropout: float = 0.3, state_size: int = 16, diff --git a/pyhealth/processors/time_image_processor.py b/pyhealth/processors/time_image_processor.py index 421998d07..a5cc4fafc 100644 --- a/pyhealth/processors/time_image_processor.py +++ b/pyhealth/processors/time_image_processor.py @@ -178,6 +178,26 @@ def _zero_image_tensor(self) -> torch.Tensor: c = 3 return torch.zeros(c, self.image_size, self.image_size) + @property + def in_channels(self) -> int: + """Channel count implied by ``mode``. + + The unified embedding sizes its patch embedding from this. Without it + the model defaulted to 3 while a greyscale task produced 1, and the + mismatch only appeared at the first forward pass: + + RuntimeError: Given groups=1, weight of size [128, 3, 16, 16], + expected input[16, 1, 224, 224] to have 3 channels + + Deriving it here means the two cannot disagree. + """ + # Must match _zero_image_tensor exactly, or a placeholder image would + # carry a different channel count from a real one. + if self.n_channels is not None: + return int(self.n_channels) + return {"1": 1, "L": 1, "LA": 2, "RGB": 3, "RGBA": 4}.get(self.mode or "RGB", 3) + + def _load_single_image(self, path: Union[str, Path]) -> torch.Tensor: """Load and transform a single image from disk. diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index cadb479ce..d892252ca 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -74,5 +74,6 @@ ClinicalNotesMIMIC4, ClinicalNotesICDLabsMIMIC4, ClinicalNotesICDLabsCXRMIMIC4, + CXRMultimodalMIMIC4, ) from .patient_linkage_mimic3 import PatientLinkageMIMIC3Task diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py index 21147303c..d7c944b5a 100644 --- a/pyhealth/tasks/multimodal_mimic4.py +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -1447,3 +1447,200 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: record["icd_codes"] = (all_icd_times, all_icd_codes) return [record] + + +class CXRMultimodalMIMIC4(BaseMultimodalMIMIC4Task): + """Admission-level MIMIC-IV mortality task with timestamped MIMIC-CXR. + + MIMIC-CXR has no ``hadm_id``. This task therefore links an image event to + an admission only when the *same subject's* ``StudyDate + StudyTime`` falls + inside that admission's observation interval. This is intentionally done + in the task (instead of a loose subject-level merge) so studies acquired + after prediction time cannot enter the sample. + + MIMIC-CXR and MIMIC-IV timestamps are both de-identified with the same + patient-specific date shift. Consequently, ``metadata.timestamp - + admission.timestamp`` is a valid hours-from-admission time just like the + lab offsets used by :meth:`_collect_labs`. + + The default is PA/AP frontal images only. Frontal views are the standard + clinical CXR representation and exclude lateral/oblique acquisitions whose + geometry makes a single lightweight patch encoder substantially less + comparable. Set ``frontal_only=False`` to retain every view. + + One sample is emitted per admission with at least one eligible CXR. This + preserves the actual CXR-to-admission linkage rather than injecting an + invalid blank image placeholder for unlinked admissions. Patient-level + splitting remains required downstream to avoid patient overlap. + + Args: + window_hours: Observation window from admission. Defaults to 24 h. + include_labs: Include the existing timestamped 10-dimensional lab + vectors in the same admission window. + include_notes: Include existing admission-context notes in the same + window. Requires the corresponding MIMIC-IV note table. + frontal_only: Keep PA/AP images only. Defaults to ``True``. + image_size: Square loader resize. Defaults to 224. + max_images: Cap images per admission after chronological sorting. The + most recent images are retained, matching ``TimeImageProcessor``. + note_source: ``"discharge"`` or ``"radiology"`` when notes are used. + """ + + task_name: str = "CXRMultimodalMIMIC4" + output_schema: Dict[str, str] = {"mortality": "binary"} + + _CXR_SCHEMA: ClassVar[Dict[str, Union[str, Tuple[str, Dict]]]] = { + # MIMIC-CXR JPEGs are grayscale. Decoding to one channel avoids a + # gratuitous 3x input expansion while VisionEmbeddingModel/PatchEmbedding + # automatically receives the processor's inferred channel count. + "cxr": ("time_image", {"image_size": 224, "mode": "L", "max_images": 4}), + } + _NOTES_SCHEMA: ClassVar[Tuple[str, Dict]] = ( + "tuple_time_text", + { + "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", + "max_length": 512, + "type_tag": "note", + }, + ) + + def __init__( + self, + window_hours: Optional[float] = 24, + include_labs: bool = False, + include_notes: bool = False, + frontal_only: bool = True, + image_size: int = 224, + max_images: Optional[int] = 4, + note_source: str = "discharge", + ) -> None: + if window_hours is not None and window_hours < 0: + raise ValueError("window_hours must be non-negative or None.") + if image_size <= 0: + raise ValueError("image_size must be positive.") + if max_images is not None and max_images <= 0: + raise ValueError("max_images must be positive or None.") + if note_source not in {"discharge", "radiology"}: + raise ValueError(f"Unsupported note_source: {note_source!r}") + + super().__init__(window_hours=window_hours) + # Included in the task cache UUID; increment on emitted-record changes. + self.cxr_pipeline_cache_version = 1 + self.include_labs = include_labs + self.include_notes = include_notes + self.frontal_only = frontal_only + self.image_size = image_size + self.max_images = max_images + self.note_source = note_source + + schema = dict(self._CXR_SCHEMA) + schema["cxr"] = ( + "time_image", + {"image_size": image_size, "mode": "L", "max_images": max_images}, + ) + if include_labs: + schema["labs"] = ("stagenet_tensor", {}) + schema["labs_mask"] = ("stagenet_tensor", {}) + if include_notes: + schema["admission_note_times"] = self._NOTES_SCHEMA + self.input_schema = schema + + suffix = ["frontal" if frontal_only else "allviews"] + if include_labs: + suffix.append("labs") + if include_notes: + suffix.append(f"notes_{note_source}") + self.task_name = "CXRMultimodalMIMIC4_" + "_".join(suffix) + + def _cxr_window_end(self, admission_time: datetime, admission: Any) -> datetime: + """Return this admission's leakage-safe CXR/lab/note cutoff.""" + dischtime = self._parse_datetime(getattr(admission, "dischtime", None)) + if dischtime is None or dischtime < admission_time: + dischtime = admission_time + if self.window_hours is None: + return dischtime + return min(dischtime, admission_time + timedelta(hours=self.window_hours)) + + def _collect_admission_cxr( + self, + patient: Any, + admission_time: datetime, + end_time: datetime, + ) -> Tuple[List[str], List[float]]: + """Collect valid CXR paths and real hours-from-admission timestamps.""" + events = patient.get_events( + event_type="metadata", start=admission_time, end=end_time + ) + paths: List[str] = [] + times: List[float] = [] + seen_dicom: set[str] = set() + for event in events: + try: + view = str(getattr(event, "viewposition", "")).upper().strip() + if self.frontal_only and view not in {"PA", "AP"}: + continue + path = str(event.image_path) + dicom_id = str(getattr(event, "dicom_id", path)) + if not path or dicom_id in seen_dicom: + continue + seen_dicom.add(dicom_id) + paths.append(path) + times.append( + self._to_hours((event.timestamp - admission_time).total_seconds()) + ) + except AttributeError: + # A malformed metadata record cannot become an image event. + continue + return paths, times + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: + if not patient.get_events(event_type="patients"): + return [] + records: List[Dict[str, Any]] = [] + for admission in patient.get_events(event_type="admissions"): + admission_time = admission.timestamp + # Discharge-note retrieval spans the admission (the section + # extraction is the temporal control), so this loop needs dischtime. + admission_dischtime = ( + self._parse_datetime(getattr(admission, "dischtime", None)) + or admission_time + ) + end_time = self._cxr_window_end(admission_time, admission) + cxr_paths, cxr_times = self._collect_admission_cxr( + patient, admission_time, end_time + ) + # Do not fabricate an image sentinel: an unreadable blank path would + # fail the image processor and, more importantly, hide cohort shift. + if not cxr_paths: + continue + + record: Dict[str, Any] = { + "patient_id": patient.patient_id, + "hadm_id": str(getattr(admission, "hadm_id", "")), + "cxr": (cxr_paths, cxr_times), + "mortality": int(getattr(admission, "hospital_expire_flag", 0) in (1, "1")), + "window_start": admission_time, + "window_end": end_time, + } + if self.include_labs: + lab_times, lab_values, lab_masks = self._collect_labs( + patient=patient, admission_time=admission_time, end_time=end_time + ) + record["labs"] = (lab_times, lab_values) + record["labs_mask"] = (lab_times, lab_masks) + if self.include_notes: + # Use the admission-context section extraction that main + # already has. A discharge summary is written with knowledge of + # the outcome, so only the admission-context sections enter a + # sample. + texts, note_times = self._collect_notes( + patient, + self.note_source, + getattr(admission, "hadm_id", None), + admission_time, + end_time=end_time, + section_headers=self.DISCHARGE_CLINICAL_HEADERS, + ) + record["admission_note_times"] = (texts, note_times) + records.append(record) + return records diff --git a/pyhealth/trainer.py b/pyhealth/trainer.py index fc264a3af..a774e0c81 100644 --- a/pyhealth/trainer.py +++ b/pyhealth/trainer.py @@ -61,6 +61,30 @@ def get_metrics_fn(mode: str) -> Callable: raise ValueError(f"Mode {mode} is not supported") +def is_text_pathway_param(name: str, text_fields, frozen_text_fields=()) -> bool: + """Whether a parameter belongs to the text pathway, for discriminative LR. + + ``encoder_lr`` exists to give a PRETRAINED encoder a gentler learning rate + than the randomly initialised layers around it, so a projection normally + keeps the base rate. + + The exception is a frozen encoder. Then every ``encoders.*`` parameter has + ``requires_grad=False``, the group is empty, and ``encoder_lr`` controls + nothing, while the only trainable text parameters, ``projections.*``, keep + the base rate. For those fields the projection IS the text pathway, so it + joins the group. + """ + frozen = set(frozen_text_fields or ()) + return any( + name.startswith(f"embedding_model.encoders.{field}.") + or ( + field in frozen + and name.startswith(f"embedding_model.projections.{field}.") + ) + for field in text_fields + ) + + class Trainer: """Trainer for PyTorch models. @@ -140,6 +164,7 @@ def train( accumulation_steps: int = 1, use_amp: bool = False, amp_dtype: str = "bf16", + encoder_lr: Optional[float] = None, ): """Trains the model. @@ -196,16 +221,49 @@ def train( # set optimizer param = list(self.model.named_parameters()) no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"] + def _decayed(n): + return not any(nd in n for nd in no_decay) + optimizer_grouped_parameters = [ { - "params": [p for n, p in param if not any(nd in n for nd in no_decay)], + "params": [p for n, p in param if _decayed(n)], "weight_decay": weight_decay, }, { - "params": [p for n, p in param if any(nd in n for nd in no_decay)], + "params": [p for n, p in param if not _decayed(n)], "weight_decay": 0.0, }, ] + if encoder_lr is not None: + # A pretrained text encoder needs a gentler rate than the randomly + # initialised layers around it. Everything else keeps the base rate + # from optimizer_params. + embedding_model = getattr(self.model, "embedding_model", None) + modality_types = getattr(embedding_model, "modality_types", {}) + text_fields = { + field + for field, modality in modality_types.items() + if getattr(modality, "value", modality) == "text" + } + if not text_fields: + raise ValueError( + "encoder_lr was set, but the model has no text encoder" + ) + frozen_text_fields = getattr(embedding_model, "_frozen_text_fields", set()) + + def _is_encoder(n): + return is_text_pathway_param(n, text_fields, frozen_text_fields) + + optimizer_grouped_parameters = [ + {"params": [p for n, p in param if _is_encoder(n) and _decayed(n)], + "weight_decay": weight_decay, "lr": encoder_lr}, + {"params": [p for n, p in param if _is_encoder(n) and not _decayed(n)], + "weight_decay": 0.0, "lr": encoder_lr}, + {"params": [p for n, p in param if not _is_encoder(n) and _decayed(n)], + "weight_decay": weight_decay}, + {"params": [p for n, p in param if not _is_encoder(n) and not _decayed(n)], + "weight_decay": 0.0}, + ] optimizer = optimizer_class(optimizer_grouped_parameters, **optimizer_params) # initialize @@ -300,10 +358,18 @@ def train( if self.exp_path is not None: self.save_ckpt(os.path.join(self.exp_path, "last.ckpt")) + # An epoch mean cannot show the difference between a run that + # starts badly and a run that becomes worse inside the epoch. + # Record the trajectory so a divergence is visible without a rerun. + _head = training_loss[: min(100, len(training_loss))] + _tail = training_loss[-min(100, len(training_loss)):] epoch_record: Dict = { "epoch": epoch, "global_step": global_step, "train_loss": sum(training_loss) / len(training_loss), + "train_loss_first_step": round(training_loss[0], 6), + "train_loss_first100": round(sum(_head) / len(_head), 6), + "train_loss_last100": round(sum(_tail) / len(_tail), 6), "epoch_time_s": round(epoch_time, 3), **{f"train_{k}": v for k, v in vram.items()}, } diff --git a/pyhealth/utils.py b/pyhealth/utils.py index efbee3977..f6cd6ff4c 100644 --- a/pyhealth/utils.py +++ b/pyhealth/utils.py @@ -1,7 +1,9 @@ +import hashlib import json import os import pickle import random +import subprocess import contextlib import numpy as np @@ -43,6 +45,82 @@ def save_json(data, filename): with open(filename, "w") as f: json.dump(data, f) + +def _git_revision(): + repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + try: + commit = subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=repo, stderr=subprocess.DEVNULL + ).decode().strip() + dirty = bool(subprocess.check_output( + ["git", "status", "--porcelain"], cwd=repo, stderr=subprocess.DEVNULL + ).decode().strip()) + return {"commit": commit, "dirty": dirty} + except Exception: + return {"commit": None, "dirty": None} + + +def _source_digest(): + """Hash the package source so code identity survives a non-git deploy. + + Cluster runs typically execute from an unpacked tarball rather than a + clone, so the git lookup returns nothing exactly where provenance matters + most. Hashing the sources keeps "which code produced this result" + answerable either way. + """ + package = os.path.dirname(os.path.abspath(__file__)) + digest = hashlib.sha256() + try: + for root, dirs, files in os.walk(package): + dirs[:] = sorted(d for d in dirs if d != "__pycache__") + for name in sorted(files): + if not name.endswith(".py"): + continue + path = os.path.join(root, name) + digest.update(os.path.relpath(path, package).encode()) + with open(path, "rb") as f: + digest.update(f.read()) + return digest.hexdigest() + except OSError: + return None + + +def _jsonable(value): + if isinstance(value, (str, int, float, bool)) or value is None: + return value + if isinstance(value, (list, tuple)): + return [_jsonable(v) for v in value] + if isinstance(value, dict): + return {str(k): _jsonable(v) for k, v in value.items()} + return str(value) + + +def write_run_config(exp_path, config): + """Persist the resolved run configuration next to the run's metrics. + + ``metrics_history.json`` records what a run scored but not the conditions + that produced it, so a frozen-encoder run and a fine-tuned one are + indistinguishable once the job's stdout is gone. Record the resolved + settings, not the raw flags, so derived conditions are recoverable. + """ + record = { + "config": {str(k): _jsonable(v) for k, v in config.items()}, + "git": _git_revision(), + "source_sha256": _source_digest(), + "torch": torch.__version__, + } + os.makedirs(exp_path, exist_ok=True) + path = os.path.join(exp_path, "run_config.json") + tmp = f"{path}.tmp.{os.getpid()}" + try: + with open(tmp, "w") as f: + json.dump(record, f, indent=2, sort_keys=True) + os.replace(tmp, path) + finally: + if os.path.exists(tmp): + os.remove(tmp) + return path + @contextlib.contextmanager def set_env(**environ): """ diff --git a/tests/test_jamba_default_depth.py b/tests/test_jamba_default_depth.py new file mode 100644 index 000000000..5951d60fe --- /dev/null +++ b/tests/test_jamba_default_depth.py @@ -0,0 +1,43 @@ +"""Proof that JambaEHR defaults to 2 transformer + 2 mamba layers. + +The published runner used library defaults of 2+6, which over-parameterised +the e2e comparison against Transformer and RNN. Both the class and the CLI +now default to 2+2. +""" + +from __future__ import annotations + +import re +import unittest +from pathlib import Path + +RUNNER = ( + Path(__file__).resolve().parents[1] + / "examples" + / "mortality_prediction" + / "unified_embedding_e2e_mimic4.py" +) + + +class TestJambaDefaultDepth(unittest.TestCase): + def test_library_defaults_to_two_plus_two(self): + from pyhealth.models.jamba_ehr import JambaEHR, JambaLayer + + self.assertEqual(JambaLayer.__init__.__defaults__[0], 2) + self.assertEqual(JambaLayer.__init__.__defaults__[1], 2) + self.assertEqual(JambaEHR.__init__.__defaults__[1], 2) + self.assertEqual(JambaEHR.__init__.__defaults__[2], 2) + layer = JambaLayer(16) + self.assertEqual(layer.schedule.count("transformer"), 2) + self.assertEqual(layer.schedule.count("mamba"), 2) + + def test_cli_defaults_to_two_plus_two(self): + src = RUNNER.read_text(encoding="utf-8") + self.assertRegex( + src, + r'--jamba-transformer-layers", type=int, default=2', + ) + self.assertRegex( + src, + r'--jamba-mamba-layers", type=int, default=2', + ) diff --git a/tests/test_run_directory_naming.py b/tests/test_run_directory_naming.py new file mode 100644 index 000000000..8d8475ab6 --- /dev/null +++ b/tests/test_run_directory_naming.py @@ -0,0 +1,72 @@ +"""Two arms of one comparison must not write to the same directory. + +The run directory was named from the model and the seed only. A paired +comparison holds both of those fixed and varies the task, so the two arms +resolved to one path and the second run overwrote the first run's +``metrics_history.json``, ``run_config.json`` and predictions CSV. The loss is +silent: the surviving directory looks like a complete run. +""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +RUNNER = ( + Path(__file__).resolve().parents[1] + / "examples" + / "mortality_prediction" + / "unified_embedding_e2e_mimic4.py" +) + + +def _exp_name_template() -> str: + """The literal assignment, read from source. + + The name is computed deep inside ``main()`` after a dataset build, so the + assignment itself is what this test pins. + """ + for line in RUNNER.read_text().splitlines(): + stripped = line.strip() + if stripped.startswith("exp_name = "): + return stripped + raise AssertionError("exp_name is no longer assigned in the runner") + + +def _render(task: str, model: str, seed: int) -> str: + template = _exp_name_template().split("=", 1)[1].strip() + args = SimpleNamespace(task=task, model=model, seed=seed) + return eval(template, {"args": args}) # noqa: S307 - a literal f-string + + +def test_the_run_directory_name_includes_the_task(): + assert "args.task" in _exp_name_template(), ( + "two arms of the same comparison would share one output directory" + ) + + +@pytest.mark.parametrize( + "left,right", + [ + (("labs_only", "transformer", 42), ("notes_labs", "transformer", 42)), + (("labs_only", "rnn", 1), ("cxr_notes_labs", "rnn", 1)), + ], +) +def test_two_arms_at_the_same_seed_get_different_directories(left, right): + assert _render(*left) != _render(*right) + + +def test_the_seed_still_separates_repeats_of_one_arm(): + assert _render("notes_labs", "transformer", 42) != _render( + "notes_labs", "transformer", 43 + ) + + +def test_the_model_still_separates_backbones(): + assert _render("notes_labs", "transformer", 42) != _render( + "notes_labs", "ehrmamba", 42 + ) diff --git a/tests/test_run_provenance_and_pathways.py b/tests/test_run_provenance_and_pathways.py new file mode 100644 index 000000000..5302e805c --- /dev/null +++ b/tests/test_run_provenance_and_pathways.py @@ -0,0 +1,315 @@ +"""Regression tests for run provenance and the two silent fallbacks. + +Each defect below made a measurement mean something other than what it said. A +run reported validation performance as test performance. A patient split became +a sample split, which leaks. A frozen-encoder run and a fine-tuned run left +identical artefacts on disk, so the condition of an earlier result could not be +recovered. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +import pytest +import torch +import torch.nn as nn + + +# ───────────────────────────────────────────────────────────────────────────── +# Run provenance +# ───────────────────────────────────────────────────────────────────────────── + + +def test_the_run_configuration_is_written_beside_the_metrics(tmp_path): + from pyhealth.utils import write_run_config + + path = write_run_config(str(tmp_path), {"task": "notes_labs", "seed": 42}) + record = json.loads(Path(path).read_text()) + + assert Path(path).name == "run_config.json" + assert record["config"]["task"] == "notes_labs" + assert record["config"]["seed"] == 42 + + +def test_code_identity_survives_a_run_from_an_unpacked_archive(tmp_path): + """A cluster run starts from a tarball, where `git rev-parse` gives nothing. + A digest of the package source identifies the code in that case. + """ + from pyhealth.utils import write_run_config + + record = json.loads(Path(write_run_config(str(tmp_path), {})).read_text()) + + assert record["source_sha256"], "no code identity was recorded" + assert len(record["source_sha256"]) >= 16 + assert "torch" in record + + +def test_the_digest_is_stable_for_unchanged_source(): + """Two runs of the same code must record the same digest, otherwise the + digest cannot show that a set of arms shared one version of the code. + """ + from pyhealth import utils + + assert utils._source_digest() == utils._source_digest() + + +def test_a_value_that_json_cannot_hold_does_not_lose_the_whole_record(tmp_path): + """A configuration holds Paths, enums and devices. If one of them raises, + the run finishes with no provenance at all, which is the case this file + exists to prevent. + """ + from pyhealth.utils import write_run_config + + config = { + "output_dir": Path("/scratch/run"), + "device": torch.device("cpu"), + "window_hours": 24, + "model": nn.Linear(2, 2), + } + record = json.loads(Path(write_run_config(str(tmp_path), config)).read_text()) + + assert set(record["config"]) == set(config) + assert record["config"]["window_hours"] == 24 + + +def test_the_temporary_file_does_not_remain_after_a_write(tmp_path): + """A partially written run_config.json is worse than none, so the write is + atomic. No temporary file may survive it. + """ + from pyhealth.utils import write_run_config + + write_run_config(str(tmp_path), {"seed": 1}) + + leftovers = [p.name for p in tmp_path.iterdir() if ".tmp." in p.name] + assert leftovers == [], f"temporary files remained: {leftovers}" + + +# ───────────────────────────────────────────────────────────────────────────── +# Discriminative learning rate for the text pathway +# ───────────────────────────────────────────────────────────────────────────── + + +TEXT_FIELDS = {"notes"} + + +def test_a_trainable_encoder_keeps_the_projection_at_the_base_rate(): + """``encoder_lr`` exists to give a PRETRAINED encoder a gentler rate. A + projection with random values must not receive that rate. + """ + from pyhealth.trainer import is_text_pathway_param + + assert is_text_pathway_param( + "embedding_model.encoders.notes.layer.0.weight", TEXT_FIELDS + ) + assert not is_text_pathway_param( + "embedding_model.projections.notes.weight", TEXT_FIELDS + ) + + +def test_a_frozen_encoder_puts_the_projection_in_the_group(): + """With the encoder frozen, every ``encoders.*`` parameter has + ``requires_grad=False``, so the group is empty and ``encoder_lr`` controls + nothing. The projection is then the only trainable text parameter, so it IS + the text pathway. + """ + from pyhealth.trainer import is_text_pathway_param + + assert is_text_pathway_param( + "embedding_model.projections.notes.weight", TEXT_FIELDS, frozen_text_fields={"notes"} + ) + + +def test_a_non_text_parameter_never_joins_the_group(): + from pyhealth.trainer import is_text_pathway_param + + for name in ( + "_unified_backbone.layers.0.weight", + "fc.weight", + "embedding_model.encoders.labs.weight", + "embedding_model.projections.labs.weight", + ): + assert not is_text_pathway_param(name, TEXT_FIELDS, frozen_text_fields={"notes"}) + + +def test_a_field_name_that_is_a_prefix_of_another_does_not_match(): + """``notes`` must not capture ``notes_extra``.""" + from pyhealth.trainer import is_text_pathway_param + + assert not is_text_pathway_param( + "embedding_model.encoders.notes_extra.weight", TEXT_FIELDS + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# Within-epoch loss trajectory +# ───────────────────────────────────────────────────────────────────────────── + + +def test_the_epoch_record_holds_the_within_epoch_trajectory(): + """An epoch mean cannot show the difference between a run that starts badly + and a run that becomes worse inside the epoch. + """ + import inspect + + from pyhealth.trainer import Trainer + + source = inspect.getsource(Trainer.train) + for field in ( + "train_loss_first_step", + "train_loss_first100", + "train_loss_last100", + ): + assert f'"{field}"' in source, f"{field} is not recorded" + + +# ───────────────────────────────────────────────────────────────────────────── +# Chest X-ray as a third modality +# ───────────────────────────────────────────────────────────────────────────── + + +CXR_ARMS = [ + ("cxr_only", {}, {"cxr"}), + ("cxr_labs", {"include_labs": True}, {"cxr", "labs", "labs_mask"}), + ( + "cxr_notes_labs", + {"include_labs": True, "include_notes": True}, + {"cxr", "labs", "labs_mask", "admission_note_times"}, + ), +] + + +@pytest.mark.parametrize("name,kwargs,expected", CXR_ARMS) +def test_each_cxr_arm_declares_the_fields_it_uses(name, kwargs, expected): + from pyhealth.tasks import CXRMultimodalMIMIC4 + + task = CXRMultimodalMIMIC4(window_hours=24, **kwargs) + + assert set(task.input_schema) == expected + + +def test_the_lab_mask_matches_the_other_tasks(): + """``StageNetTensorProcessor`` takes no arguments. A processor option here + that the other tasks do not use would raise at dataset build time. + """ + from pyhealth.processors import StageNetTensorProcessor + from pyhealth.tasks import CXRMultimodalMIMIC4 + + task = CXRMultimodalMIMIC4(window_hours=24, include_labs=True) + name, options = task.input_schema["labs_mask"] + + assert name == "stagenet_tensor" + assert options == {} + StageNetTensorProcessor(**options) # must not raise + + +def test_an_image_uses_the_same_time_convention_as_a_laboratory_value(): + """``StudyDate`` and ``StudyTime`` give hours from admission, which is what + the unified embedding reads for every modality. + """ + from pyhealth.tasks import CXRMultimodalMIMIC4 + + task = CXRMultimodalMIMIC4(window_hours=24) + assert task.input_schema["cxr"][0] == "time_image" + + +def test_sunlab_accepts_resized_images_as_well_as_images(tmp_path): + """The resized set holds flattened ``{dicom_id}.jpg`` files under + ``resized_images``. Requiring a directory literally named ``images`` failed + on a complete 377,110-image cohort. + """ + import pandas as pd + from pyhealth.datasets.mimic4 import MIMIC4CXRSunlabDataset + + root = tmp_path / "cxr" + (root / "resized_images").mkdir(parents=True) + pd.DataFrame( + {"dicom_id": ["abc123"], "StudyTime": ["123045.0"]} + ).to_csv(root / "mimic-cxr-2.0.0-metadata.csv", index=False) + + MIMIC4CXRSunlabDataset.prepare_metadata( + object.__new__(MIMIC4CXRSunlabDataset), str(root) + ) + + out = pd.read_csv(root / "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv") + path = str(out["image_path"].iloc[0]) + assert path.endswith(f"resized_images{os.sep}abc123.jpg") or path.endswith( + "resized_images/abc123.jpg" + ) + assert str(out["studytime"].iloc[0]).zfill(6) == "123045" + + +def test_sunlab_writes_metadata_to_cache_when_root_is_unwritable(tmp_path): + """PhysioNet roots are typically read-only. Writing the derived CSV there + raised PermissionError after a complete cohort had already been found. + """ + import pandas as pd + from pyhealth.datasets.mimic4 import MIMIC4CXRSunlabDataset + + root = tmp_path / "cxr" + (root / "resized_images").mkdir(parents=True) + pd.DataFrame( + {"dicom_id": ["abc"], "StudyTime": ["93000"], "subject_id": ["1"]} + ).to_csv(root / "mimic-cxr-2.0.0-metadata.csv", index=False) + cache = tmp_path / "cache" + cache.mkdir() + os.chmod(root, 0o555) + try: + dest = MIMIC4CXRSunlabDataset.prepare_metadata( + object.__new__(MIMIC4CXRSunlabDataset), + str(root), + cache_dir=str(cache), + ) + finally: + os.chmod(root, 0o755) + + assert dest.startswith(str(cache)) + assert Path(dest).is_file() + assert not (root / "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv").exists() + written = pd.read_csv(dest) + path = str(written.loc[0, "image_path"]) + assert path.endswith(f"resized_images{os.sep}abc.jpg") or path.endswith( + "resized_images/abc.jpg" + ) + + +def test_sunlab_reports_both_directory_names_when_neither_exists(tmp_path): + import pandas as pd + from pyhealth.datasets.mimic4 import MIMIC4CXRSunlabDataset + + root = tmp_path / "cxr" + root.mkdir() + pd.DataFrame( + {"dicom_id": ["abc123"], "StudyTime": ["1"]} + ).to_csv(root / "mimic-cxr-2.0.0-metadata.csv", index=False) + + with pytest.raises(FileNotFoundError, match="resized_images"): + MIMIC4CXRSunlabDataset.prepare_metadata( + object.__new__(MIMIC4CXRSunlabDataset), str(root) + ) + + + +def test_the_image_channel_count_comes_from_the_processor(): + """The unified embedding sizes its patch embedding from + ``processor.in_channels``. Without that attribute it fell back to 3 while a + greyscale CXR task produced 1, and the mismatch surfaced only at the first + forward pass: + + RuntimeError: Given groups=1, weight of size [128, 3, 16, 16], + expected input[16, 1, 224, 224] to have 3 channels + """ + from pyhealth.processors import TimeImageProcessor + + assert TimeImageProcessor(image_size=224, mode="L").in_channels == 1 + assert TimeImageProcessor(image_size=224, mode="RGB").in_channels == 3 + + +def test_a_placeholder_image_has_the_same_channels_as_a_real_one(): + from pyhealth.processors import TimeImageProcessor + + for mode in ("L", "RGB"): + processor = TimeImageProcessor(image_size=64, mode=mode) + assert processor._zero_image_tensor().shape[0] == processor.in_channels