From bc2e1a189fd5c81f3e70f9a86f9843a492403997 Mon Sep 17 00:00:00 2001 From: Rian354 Date: Mon, 3 Aug 2026 22:22:17 -0400 Subject: [PATCH] Multimodal SSL pretraining: MAE, SimMIM, I-JEPA and V-JEPA with multi-arch backbones Add pyhealth/models/pretrain/, a self-supervised pretraining stack for the multimodal clinical-sequence encoder. build_backbone selects among transformer, jamba and mamba backbones behind a uniform (x, mask) -> (emb, cls) contract, so every objective composes with every backbone. The mamba layer pads masked positions so its causal conv cannot leak across the mask, and RoPE is available for long-sequence extrapolation. Four objectives share the unified embedding and one mask generator: MAE (decoder reconstruction), SimMIM (masked-token linear head), and I-JEPA / V-JEPA (EMA-target latent prediction, with V-JEPA adding multi-scale target blocks). PretrainTrainer is DDP-native, supports AMP and gradient accumulation, drives the EMA target encoder for the JEPA variants, and skips DistributedSampler for IterableDataset streams that already shard internally. It resumes atomically from last.ckpt plus _resume.pt, and warm-resumes from weights alone when optimizer state is missing, so preempted cluster jobs continue instead of restarting. Per-modality losses are tracked to metrics_history.json and W&B. scripts/pretrain_ssl.py drives it (--arch / --method / --task, standardized 128/2/4 encoders, notes_only via include_labs=False). run_full_pretrain.py replays tuned hyperparameters from a best_params JSON and can override them with --extra. SLURM and Condor runners skip only encoders that reached a full 50 epochs, so partial runs resume, and set expandable_segments to avoid fragmentation OOM. W&B logging is env-gated, namespaced and auto-tagged by arch/method/task. Includes unit tests for masking, all four objectives, backbone contracts, RoPE, and checkpoint transfer into the downstream Transformer. --- configs/pretrain/base.yaml | 70 +++ configs/pretrain/ijepa_labs_only.yaml | 7 + configs/pretrain/ijepa_notes_labs.yaml | 7 + configs/pretrain/mae_labs_only.yaml | 7 + configs/pretrain/mae_notes_labs.yaml | 8 + configs/pretrain/simmim_labs_only.yaml | 7 + configs/pretrain/simmim_notes_labs.yaml | 8 + configs/pretrain/vjepa_labs_only.yaml | 9 + configs/pretrain/vjepa_notes_labs.yaml | 11 + pyhealth/_wandb.py | 116 +++++ pyhealth/models/pretrain/__init__.py | 35 ++ pyhealth/models/pretrain/backbones.py | 150 +++++++ pyhealth/models/pretrain/jepa.py | 556 +++++++++++++++++++++++ pyhealth/models/pretrain/mae.py | 295 ++++++++++++ pyhealth/models/pretrain/masking.py | 183 ++++++++ pyhealth/models/pretrain/rope.py | 266 +++++++++++ pyhealth/models/pretrain/simmim.py | 181 ++++++++ pyhealth/models/pretrain/trainer.py | 543 ++++++++++++++++++++++ pyhealth/models/pretrain/utils.py | 47 ++ scripts/kill_ddp.sh | 15 + scripts/pretrain_ssl.py | 531 ++++++++++++++++++++++ scripts/run_full_pretrain.py | 102 +++++ scripts/run_full_pretrain_local.sh | 47 ++ scripts/run_fullpt_condor.sh | 8 + scripts/slurm/full_pretrain_cc.sh | 49 ++ tests/test_pretrain.py | 574 ++++++++++++++++++++++++ tests/test_pretrain_backbones.py | 106 +++++ 27 files changed, 3938 insertions(+) create mode 100644 configs/pretrain/base.yaml create mode 100644 configs/pretrain/ijepa_labs_only.yaml create mode 100644 configs/pretrain/ijepa_notes_labs.yaml create mode 100644 configs/pretrain/mae_labs_only.yaml create mode 100644 configs/pretrain/mae_notes_labs.yaml create mode 100644 configs/pretrain/simmim_labs_only.yaml create mode 100644 configs/pretrain/simmim_notes_labs.yaml create mode 100644 configs/pretrain/vjepa_labs_only.yaml create mode 100644 configs/pretrain/vjepa_notes_labs.yaml create mode 100644 pyhealth/_wandb.py create mode 100644 pyhealth/models/pretrain/__init__.py create mode 100644 pyhealth/models/pretrain/backbones.py create mode 100644 pyhealth/models/pretrain/jepa.py create mode 100644 pyhealth/models/pretrain/mae.py create mode 100644 pyhealth/models/pretrain/masking.py create mode 100644 pyhealth/models/pretrain/rope.py create mode 100644 pyhealth/models/pretrain/simmim.py create mode 100644 pyhealth/models/pretrain/trainer.py create mode 100644 pyhealth/models/pretrain/utils.py create mode 100644 scripts/kill_ddp.sh create mode 100644 scripts/pretrain_ssl.py create mode 100644 scripts/run_full_pretrain.py create mode 100644 scripts/run_full_pretrain_local.sh create mode 100644 scripts/run_fullpt_condor.sh create mode 100644 scripts/slurm/full_pretrain_cc.sh create mode 100644 tests/test_pretrain.py create mode 100644 tests/test_pretrain_backbones.py diff --git a/configs/pretrain/base.yaml b/configs/pretrain/base.yaml new file mode 100644 index 000000000..20c309c3d --- /dev/null +++ b/configs/pretrain/base.yaml @@ -0,0 +1,70 @@ +# Base config for SSL pretraining on unified multimodal sequences. +# Use with scripts/pretrain_ssl.py. + +method: mae + +# Encoder backbone: transformer | jamba | mamba. +arch: transformer + +# Model — standardized compute: 128-dim, 2 layers, 4 heads (head_dim 32), +# matching the downstream e2e backbone so SSL encoders transfer 1:1. Bumped up +# from 64/1 (too narrow: V-JEPA drifted, downstream sat near prevalence); VRAM is +# BERT-bound so the larger backbone is effectively free on 48 GB RTX 6000 Ada. +embedding_dim: 128 +heads: 4 +num_layers: 2 +dropout: 0.1 + +# Mamba / Jamba backbone knobs (ignored by the transformer backbone). +state_size: 16 +conv_kernel: 4 +jamba_transformer_layers: 1 +jamba_mamba_layers: 1 + +# MAE decoder (ignored for SimMIM / I-JEPA). Discarded after pretraining, so it +# does NOT count toward the deployed-encoder 128/2 standard — kept at 2 layers so +# the reconstruction objective remains learnable. +decoder_layers: 2 +decoder_heads: 2 +decoder_dim: 128 +norm_pix_loss: false + +# I-JEPA / V-JEPA predictor (ignored for MAE / SimMIM). Also discarded after +# pretraining; 2 layers keeps latent prediction non-degenerate at 128-dim. +predictor_layers: 2 +predictor_heads: 2 +predictor_dim: 128 +ema_decay: 0.996 +ema_end: 1.0 +num_target_blocks: 4 +target_block_len: 4 + +# Masking +mask_ratio: 0.5 +mask_strategy: block +lab_mask_ratio: null +text_mask_ratio: null + +# Training +epochs: 50 +batch_size: 32 +lr: 0.0001 +weight_decay: 0.05 +max_grad_norm: 1.0 +scheduler: cosine +warmup_steps: 1000 +save_every_n_epochs: 5 + +# Task +observation_window_hours: 24 +note_extraction: regex +note_source: discharge +icd_codes: false +include_vitals: false +tokenizer_model: null +freeze_encoder: false +text_finetune_mode: full + +# Data / compute +num_workers: 4 +seed: 42 diff --git a/configs/pretrain/ijepa_labs_only.yaml b/configs/pretrain/ijepa_labs_only.yaml new file mode 100644 index 000000000..f11ab1e02 --- /dev/null +++ b/configs/pretrain/ijepa_labs_only.yaml @@ -0,0 +1,7 @@ +_inherit: base.yaml + +method: ijepa +task: labs_only +num_target_blocks: 4 +target_block_len: 4 +output_dir: ./output/pretrain/ijepa_labs_only diff --git a/configs/pretrain/ijepa_notes_labs.yaml b/configs/pretrain/ijepa_notes_labs.yaml new file mode 100644 index 000000000..1b9914c32 --- /dev/null +++ b/configs/pretrain/ijepa_notes_labs.yaml @@ -0,0 +1,7 @@ +_inherit: base.yaml + +method: ijepa +task: notes_labs +num_target_blocks: 6 +target_block_len: 3 +output_dir: ./output/pretrain/ijepa_notes_labs diff --git a/configs/pretrain/mae_labs_only.yaml b/configs/pretrain/mae_labs_only.yaml new file mode 100644 index 000000000..81c118521 --- /dev/null +++ b/configs/pretrain/mae_labs_only.yaml @@ -0,0 +1,7 @@ +_inherit: base.yaml + +method: mae +task: labs_only +mask_ratio: 0.5 +lab_mask_ratio: 0.5 +output_dir: ./output/pretrain/mae_labs_only diff --git a/configs/pretrain/mae_notes_labs.yaml b/configs/pretrain/mae_notes_labs.yaml new file mode 100644 index 000000000..c77773960 --- /dev/null +++ b/configs/pretrain/mae_notes_labs.yaml @@ -0,0 +1,8 @@ +_inherit: base.yaml + +method: mae +task: notes_labs +mask_ratio: 0.5 +lab_mask_ratio: 0.5 +text_mask_ratio: 0.25 +output_dir: ./output/pretrain/mae_notes_labs diff --git a/configs/pretrain/simmim_labs_only.yaml b/configs/pretrain/simmim_labs_only.yaml new file mode 100644 index 000000000..688cc4e13 --- /dev/null +++ b/configs/pretrain/simmim_labs_only.yaml @@ -0,0 +1,7 @@ +_inherit: base.yaml + +method: simmim +task: labs_only +mask_ratio: 0.5 +lab_mask_ratio: 0.5 +output_dir: ./output/pretrain/simmim_labs_only diff --git a/configs/pretrain/simmim_notes_labs.yaml b/configs/pretrain/simmim_notes_labs.yaml new file mode 100644 index 000000000..a49bdda8d --- /dev/null +++ b/configs/pretrain/simmim_notes_labs.yaml @@ -0,0 +1,8 @@ +_inherit: base.yaml + +method: simmim +task: notes_labs +mask_ratio: 0.5 +lab_mask_ratio: 0.5 +text_mask_ratio: 0.25 +output_dir: ./output/pretrain/simmim_notes_labs diff --git a/configs/pretrain/vjepa_labs_only.yaml b/configs/pretrain/vjepa_labs_only.yaml new file mode 100644 index 000000000..a1f806704 --- /dev/null +++ b/configs/pretrain/vjepa_labs_only.yaml @@ -0,0 +1,9 @@ +_inherit: base.yaml + +method: vjepa +task: labs_only +num_target_blocks: 4 +target_block_scales: [2, 4, 8] +require_multimodal_blocks: false +normalize_targets: true +output_dir: ./output/pretrain/vjepa_labs_only diff --git a/configs/pretrain/vjepa_notes_labs.yaml b/configs/pretrain/vjepa_notes_labs.yaml new file mode 100644 index 000000000..81e811cf9 --- /dev/null +++ b/configs/pretrain/vjepa_notes_labs.yaml @@ -0,0 +1,11 @@ +_inherit: base.yaml + +method: vjepa +task: notes_labs +# More, multi-scale blocks for the longer multimodal sequence; prefer windows +# that span both labs and notes so the model learns cross-modal prediction. +num_target_blocks: 6 +target_block_scales: [2, 4, 8, 16] +require_multimodal_blocks: true +normalize_targets: true +output_dir: ./output/pretrain/vjepa_notes_labs diff --git a/pyhealth/_wandb.py b/pyhealth/_wandb.py new file mode 100644 index 000000000..d5507e691 --- /dev/null +++ b/pyhealth/_wandb.py @@ -0,0 +1,116 @@ +"""Thin, opt-in Weights & Biases logging helper. + +Everything here is a no-op unless ``WANDB_PROJECT`` is set in the environment +(and ``wandb`` is importable), so importing/using it never breaks a run that +doesn't want tracking. Entity/project/mode come from the standard W&B env vars: + + WANDB_PROJECT e.g. pyhealth-multimodal (required to enable logging) + WANDB_ENTITY e.g. pyhealth-multimodal (team; optional) + WANDB_MODE online (default) | offline (offline for no-internet nodes; + sync later with `wandb sync`) + +Failures (no network, bad key, wandb missing) degrade to a warning + no-op — +tracking must never take down training. +""" +from __future__ import annotations + +import os +from typing import Any, Dict, Optional + + +def enabled() -> bool: + return bool(os.environ.get("WANDB_PROJECT")) + + +def sweeps_project() -> Optional[str]: + """Dedicated project for Optuna sweeps, kept separate from the real experiment + runs so the main project isn't cluttered by per-trial study runs. Override with + WANDB_PROJECT_SWEEPS; otherwise it's ``-sweeps``.""" + p = os.environ.get("WANDB_PROJECT") + return os.environ.get("WANDB_PROJECT_SWEEPS") or (f"{p}-sweeps" if p else None) + + +def init_run(config: Optional[Dict[str, Any]] = None, name: Optional[str] = None, + group: Optional[str] = None, job_type: Optional[str] = None, + project: Optional[str] = None, tags: Optional[list] = None): + """Start a W&B run, or return None if disabled/unavailable.""" + if not enabled(): + return None + try: + import wandb + except ImportError: + print("[wandb] WANDB_PROJECT set but wandb not installed — skipping tracking.") + return None + try: + return wandb.init( + project=project or os.environ["WANDB_PROJECT"], + entity=os.environ.get("WANDB_ENTITY"), + name=name, + group=group, + job_type=job_type, + config=config or {}, + tags=[t for t in (tags or []) if t], + reinit=True, + ) + except Exception as e: # network/auth/etc. — never fail the run over telemetry + print(f"[wandb] init failed ({e}) — continuing without tracking.") + return None + + +# Canonical metric namespaces so W&B groups panels into a few tidy sections +# (val/, test/, best/, loss/, sys/) instead of dozens of flat, redundant keys. +# Applied at log()/summary() so every caller stays consistent for free. Only +# affects W&B display names; on-disk metrics_history.json keeps its raw keys. +_CANON = { + "pr_auc": "val/pr_auc", "val_pr_auc": "val/pr_auc", "val_roc_auc": "val/roc_auc", + "val_f1": "val/f1", "val_accuracy": "val/accuracy", "val_loss": "val/loss", + "test_pr_auc": "test/pr_auc", "test_roc_auc": "test/roc_auc", "test_f1": "test/f1", + "test_accuracy": "test/accuracy", "test_loss": "test/loss", "test_test_loss": "test/loss", + "test_n": "test/n", "test_pos": "test/pos", + "best_pr_auc": "best/pr_auc", "best_val_pr_auc": "best/pr_auc", "best_epoch": "best/epoch", + "total": "loss/total", "train_loss": "loss/train", + "train_vram_peak_mb": "sys/vram_peak_mb", "train_vram_allocated_mb": "sys/vram_allocated_mb", + "learning_rate": "sys/lr", "epoch_time_s": "sys/epoch_time_s", "global_step": "sys/global_step", +} + + +def _canon_key(k: str) -> str: + if k in _CANON: + return _CANON[k] + if k.startswith("modality_") or k.startswith("scale_"): # per-modality / multi-scale SSL loss + return "loss/" + k + if k.startswith("train_loss_"): # per-component supervised loss + return "loss/" + k[len("train_loss_"):] + return k # already-namespaced (hp/, arch/, sweep/, epoch, trial, ...) pass through + + +def _canon(record: Dict[str, Any]) -> Dict[str, Any]: + return {_canon_key(k): v for k, v in record.items()} + + +def log(run, record: Dict[str, Any], step: Optional[int] = None) -> None: + if run is None: + return + try: + run.log(_canon(dict(record)), step=step) + except Exception: + pass + + +def summary(run, record: Dict[str, Any]) -> None: + if run is None: + return + try: + for k, v in _canon(record).items(): + run.summary[k] = v + except Exception: + pass + + +def finish(run) -> None: + if run is None: + return + try: + run.finish() + except Exception: + pass diff --git a/pyhealth/models/pretrain/__init__.py b/pyhealth/models/pretrain/__init__.py new file mode 100644 index 000000000..4e63db61e --- /dev/null +++ b/pyhealth/models/pretrain/__init__.py @@ -0,0 +1,35 @@ +"""Self-supervised pretraining models for PyHealth multimodal sequences. + +Models: +- :class:`MultimodalMaskedAutoencoder` — true MAE with a transformer decoder. +- :class:`MultimodalSimMIM` — SimMIM-style linear-head reconstruction. +- :class:`MultimodalIJEPA` — I-JEPA latent predictive architecture. +- :class:`MultimodalVJEPA` — V-JEPA: multi-scale spans + location/scale-aware + per-position latent prediction (extends I-JEPA). + +Utilities: +- :class:`UnifiedMaskGenerator` — masking strategies for unified event seqs. +- :func:`apply_mask_token` — replace masked positions with a learnable token. +""" + +from .jepa import MultimodalIJEPA, MultimodalVJEPA +from .mae import MultimodalMaskedAutoencoder, PerModalityMAEDecoder +from .masking import UnifiedMaskGenerator, apply_mask_token, random_mask_like +from .rope import RoPEMultiHeadedAttention, RoPETransformerLayer, RotaryPositionEmbedding +from .simmim import MultimodalSimMIM +from .trainer import PretrainTrainer + +__all__ = [ + "MultimodalMaskedAutoencoder", + "MultimodalSimMIM", + "MultimodalIJEPA", + "MultimodalVJEPA", + "PerModalityMAEDecoder", + "UnifiedMaskGenerator", + "apply_mask_token", + "random_mask_like", + "RotaryPositionEmbedding", + "RoPEMultiHeadedAttention", + "RoPETransformerLayer", + "PretrainTrainer", +] diff --git a/pyhealth/models/pretrain/backbones.py b/pyhealth/models/pretrain/backbones.py new file mode 100644 index 000000000..75ba18ef2 --- /dev/null +++ b/pyhealth/models/pretrain/backbones.py @@ -0,0 +1,150 @@ +"""Sequence-encoder backbones for SSL pretraining. + +Every SSL method (MAE / SimMIM / I-JEPA / V-JEPA) injects a ``backbone`` +(a.k.a. ``context_encoder``) that must satisfy the same contract as +:class:`pyhealth.models.transformer.TransformerLayer`:: + + emb, cls = backbone(x, mask) + # x: (B, S, E) input sequence + # mask:(B, S) 1 = valid, 0 = pad (may be float or bool) + # emb: (B, S, E) per-step encoded features + # cls: (B, E) pooled vector (unused by the SSL methods, kept for parity) + +``build_backbone`` is the single place the training scripts and the Optuna +sweeps construct an encoder, so a new architecture only has to be added here. + +Supported ``arch`` values: + - ``"transformer"`` -> :class:`TransformerLayer` (or RoPE variant if ``use_rope``) + - ``"jamba"`` -> :class:`JambaLayer` (interleaved attention + Mamba) + - ``"mamba"`` -> :class:`MambaLayer` (this module; stacks ``MambaBlock``) + +The Jamba layer already matches the contract; ``TransformerLayer`` does too. +``MambaBlock`` is a single ``forward(x) -> x`` residual block, so this module +wraps a stack of them into a mask-aware layer that returns ``(emb, cls)``. +""" +from __future__ import annotations + +from typing import Optional, Tuple + +import torch +import torch.nn as nn + +from pyhealth.models.ehrmamba import MambaBlock, RMSNorm +from pyhealth.models.jamba_ehr import JambaLayer +from pyhealth.models.transformer import TransformerLayer +from pyhealth.models.utils import get_last_visit + +__all__ = ["MambaLayer", "build_backbone", "ARCH_CHOICES"] + +ARCH_CHOICES = ("transformer", "jamba", "mamba") + + +class MambaLayer(nn.Module): + """A stack of :class:`MambaBlock` layers exposing the standard backbone + contract ``forward(x, mask) -> (emb, cls)``. + + Padded positions are zeroed on input. Because ``MambaBlock`` is causal + (left-padded conv + left-to-right SSM scan) and pad tokens sit at the end + of each sequence, they cannot leak into earlier valid positions, so a + single input masking is sufficient. + + Args: + feature_size: hidden/embedding dimension ``E``. + num_layers: number of stacked Mamba blocks. Default 2. + dropout: dropout on the output features. Default 0.0. + state_size: SSM state size per channel. Default 16. + conv_kernel: causal conv kernel size inside each block. Default 4. + """ + + def __init__( + self, + feature_size: int, + num_layers: int = 2, + dropout: float = 0.0, + state_size: int = 16, + conv_kernel: int = 4, + ): + super().__init__() + self.blocks = nn.ModuleList( + [ + MambaBlock(d_model=feature_size, state_size=state_size, conv_kernel=conv_kernel) + for _ in range(num_layers) + ] + ) + self.norm = RMSNorm(feature_size) + self.dropout = nn.Dropout(dropout) + + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None, register_hook: bool = False + ) -> Tuple[torch.Tensor, torch.Tensor]: + if mask is not None: + x = x * mask.unsqueeze(-1).to(x.dtype) + for block in self.blocks: + x = block(x) + x = self.norm(x) + emb = self.dropout(x) + cls = get_last_visit(emb, mask) if mask is not None else emb[:, -1, :] + return emb, cls + + +def build_backbone( + arch: str, + feature_size: int, + num_layers: int = 2, + heads: int = 4, + dropout: float = 0.1, + *, + # transformer + use_rope: bool = False, + rope_max_seq_len: int = 8192, + rope_base: float = 10000.0, + rope_scaling: float = 1.0, + # mamba / jamba + state_size: int = 16, + conv_kernel: int = 4, + # jamba layer mix (num_layers is ignored for jamba) + num_transformer_layers: int = 1, + num_mamba_layers: int = 1, +) -> nn.Module: + """Construct an SSL encoder backbone satisfying the standard contract. + + ``arch`` is one of :data:`ARCH_CHOICES`. ``feature_size`` (== embedding + dim), ``num_layers`` and ``heads`` are the standardized-size knobs; the + remaining kwargs are arch-specific and ignored where not applicable. + """ + arch = arch.lower() + if arch == "transformer": + if use_rope: + from pyhealth.models.pretrain.rope import RoPETransformerLayer + + return RoPETransformerLayer( + feature_size=feature_size, + heads=heads, + dropout=dropout, + num_layers=num_layers, + rope_max_seq_len=rope_max_seq_len, + rope_base=rope_base, + rope_scaling=rope_scaling, + ) + return TransformerLayer( + feature_size=feature_size, heads=heads, dropout=dropout, num_layers=num_layers + ) + if arch == "mamba": + return MambaLayer( + feature_size=feature_size, + num_layers=num_layers, + dropout=dropout, + state_size=state_size, + conv_kernel=conv_kernel, + ) + if arch == "jamba": + return JambaLayer( + feature_size=feature_size, + num_transformer_layers=num_transformer_layers, + num_mamba_layers=num_mamba_layers, + heads=heads, + dropout=dropout, + state_size=state_size, + conv_kernel=conv_kernel, + ) + raise ValueError(f"Unknown backbone arch '{arch}'. Choices: {ARCH_CHOICES}") diff --git a/pyhealth/models/pretrain/jepa.py b/pyhealth/models/pretrain/jepa.py new file mode 100644 index 000000000..0b146bf30 --- /dev/null +++ b/pyhealth/models/pretrain/jepa.py @@ -0,0 +1,556 @@ +"""I-JEPA / V-JEPA pretraining for unified multimodal event sequences. + +A context encoder processes unmasked context positions, a target encoder +(processing the full input) provides regression targets via an EMA update, and +a small predictor network maps context representations to the EMA-target latents +of masked "target blocks". + +Because the objective operates purely in representation space, the model does +not reconstruct noisy raw values. This makes it attractive for multimodal EHR +data where labs, codes, and text have very different noise characteristics. + +The predictor is **location aware**: every target position receives its own +query (``base_query + sinusoidal_pos_emb(position)`` and, for V-JEPA, a +``scale_embed``), so the predictor can produce a distinct prediction per target +position. Prediction and loss are computed in a fully batched, per-position way +(``pred[target] vs target_latent[target]``), so there is no block-splitting and +no all-pairs broadcasting. + +References: + Mahmoud Assran et al., "Self-Supervised Learning from Images with a Joint + Embedding Predictive Architecture", CVPR 2023. + Adrien Bardes et al., "Revisiting Feature Prediction for Learning Visual + Representations from Video" (V-JEPA), 2024. +""" + +from __future__ import annotations + +import copy +import math +from typing import Any, Optional + +import torch +import torch.nn as nn + +from ..embedding.unified import UnifiedMultimodalEmbeddingModel +from ..transformer import TransformerLayer + + +class MultimodalIJEPA(nn.Module): + """I-JEPA pretrainer for unified multimodal event sequences. + + Args: + embedding_model: Unified embedding model. + context_encoder: Sequence encoder backbone. + predictor_layers: Number of layers in the predictor network. + predictor_heads: Attention heads in the predictor. + predictor_dim: Hidden dim of the predictor. + target_ema_decay: EMA momentum for the target encoder. + Default 0.996, increased to ``target_ema_end`` over training via + ``set_ema_decay`` (the trainer wires this automatically). + target_ema_end: Final EMA momentum (cosine schedule). Default 1.0. + num_target_blocks: Number of contiguous target blocks to predict per + sample. Default 4. + target_block_len: Length of each target block. Default 4. + min_context_len: Minimum number of visible context positions required. + Default 4. + normalize_targets: LayerNorm (no affine) the EMA targets before the + loss (standard JEPA practice; reduces representation collapse). + + Inputs: + Same dict as ``UnifiedMultimodalEmbeddingModel.forward``. + + Outputs: + Dict with ``loss``, ``loss_dict``, ``context_pred`` (per-position + predictions for target positions), ``target_embs`` (detached EMA + targets), ``target_mask``, ``context_mask``, ``event_mask``, + ``type_ids``. + """ + + def __init__( + self, + embedding_model: UnifiedMultimodalEmbeddingModel, + context_encoder: nn.Module, + predictor_layers: int = 6, + predictor_heads: int = 8, + predictor_dim: Optional[int] = None, + target_ema_decay: float = 0.996, + target_ema_end: float = 1.0, + num_target_blocks: int = 4, + target_block_len: int = 4, + min_context_len: int = 4, + normalize_targets: bool = True, + ): + super().__init__() + self.embedding_model = embedding_model + self.context_encoder = context_encoder + self.embedding_dim = embedding_model.embedding_dim + self.predictor_dim = predictor_dim or self.embedding_dim + + # Target encoder is an EMA copy of the context encoder. + self.target_encoder = copy.deepcopy(context_encoder) + for p in self.target_encoder.parameters(): + p.requires_grad = False + + # Predictor: small Transformer that consumes context latents at context + # positions and location-aware queries at target positions. + self.predictor_context_proj = ( + nn.Linear(self.embedding_dim, self.predictor_dim) + if self.predictor_dim != self.embedding_dim + else nn.Identity() + ) + self.predictor_query = nn.Parameter(torch.randn(1, 1, self.predictor_dim)) + nn.init.trunc_normal_(self.predictor_query, std=0.02) + self.predictor = TransformerLayer( + feature_size=self.predictor_dim, + heads=predictor_heads, + dropout=0.0, + num_layers=predictor_layers, + ) + self.predictor_norm = nn.LayerNorm(self.predictor_dim) + self.predictor_out_proj = nn.Linear(self.predictor_dim, self.embedding_dim) + + self.normalize_targets = normalize_targets + self.target_norm = ( + nn.LayerNorm(self.embedding_dim, elementwise_affine=False) + if normalize_targets + else nn.Identity() + ) + + self.target_ema_decay = target_ema_decay + self.target_ema_end = target_ema_end + self.num_target_blocks = num_target_blocks + self.target_block_len = target_block_len + self.min_context_len = min_context_len + self._ema_start = target_ema_decay + self._global_step = 0 + + def set_ema_decay(self, step: int, total_steps: int) -> None: + """Cosine schedule from ``target_ema_decay`` to ``target_ema_end``.""" + progress = min(1.0, step / max(1, total_steps)) + self._global_step = step + self.target_ema_decay = ( + self._ema_start + + (self.target_ema_end - self._ema_start) + * (1 - math.cos(progress * math.pi)) + / 2 + ) + + @torch.no_grad() + def update_target_encoder(self) -> None: + """EMA update of target encoder from context encoder.""" + m = self.target_ema_decay + for param_q, param_k in zip( + self.context_encoder.parameters(), self.target_encoder.parameters() + ): + param_k.data.mul_(m).add_((1.0 - m) * param_q.detach().data) + + @staticmethod + def _sinusoidal_pos_emb( + positions: torch.Tensor, dim: int, device: torch.device + ) -> torch.Tensor: + """Parameter-free sinusoidal embedding of integer positions. + + Args: + positions: ``(S,)`` long tensor of position indices. + dim: Output dimension. + + Returns: + ``(S, dim)`` float tensor. + """ + half = dim // 2 + freqs = torch.exp( + torch.arange(half, device=device, dtype=torch.float32) + * (-math.log(10000.0) / max(1, half - 1)) + ) + ang = positions.float().unsqueeze(-1) * freqs # (S, half) + emb = torch.cat([torch.sin(ang), torch.cos(ang)], dim=-1) + if emb.shape[-1] < dim: # pad odd dims + emb = torch.cat( + [emb, torch.zeros(emb.shape[0], dim - emb.shape[-1], device=device)], + dim=-1, + ) + return emb + + def _sample_target_blocks( + self, + event_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Sample non-overlapping single-scale target blocks + context. + + Returns: + target_mask: ``(B, S)`` bool, True = target positions. + context_mask: ``(B, S)`` bool, True = context positions. + """ + B, S = event_mask.shape + device = event_mask.device + valid = event_mask.bool() + target_mask = torch.zeros_like(event_mask, dtype=torch.bool) + + for b in range(B): + valid_positions = valid[b].nonzero(as_tuple=False).flatten() + n_valid = int(valid_positions.numel()) + if n_valid < self.min_context_len + self.target_block_len: + # Too short: mask a small random subset, leaving context. + n_target = min( + max(0, n_valid - self.min_context_len), + self.num_target_blocks * self.target_block_len, + ) + if n_target > 0: + perm = torch.randperm(n_valid, device=device) + target_mask[b, valid_positions[perm[:n_target]]] = True + continue + + max_targets = n_valid - self.min_context_len + used = torch.zeros(n_valid, dtype=torch.bool, device=device) + n_blocks = 0 + attempts = 0 + max_attempts = self.num_target_blocks * 20 + while n_blocks < self.num_target_blocks and attempts < max_attempts: + attempts += 1 + block_len = min(self.target_block_len, max_targets) + if block_len <= 0: + break + start = int( + torch.randint(0, n_valid - block_len + 1, (1,), device=device).item() + ) + span = torch.arange(start, start + block_len, device=device) + if used[span].any(): + continue + if int(used.sum().item()) + block_len > max_targets: + continue + used[span] = True + target_mask[b, valid_positions[span]] = True + n_blocks += 1 + + context_mask = valid & ~target_mask + return target_mask, context_mask + + def _encode_and_predict( + self, + sequence: torch.Tensor, + event_mask: torch.Tensor, + target_mask: torch.Tensor, + context_mask: torch.Tensor, + scale_ids: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Shared context/target encode + location-aware batched prediction. + + Returns ``(pred, target_latent)`` both of shape ``(B, S, E)``. + """ + B, S, E = sequence.shape + device = sequence.device + + # Context encoder: attend ONLY over context positions (no target leak). + context_input = sequence.clone() + context_input[~context_mask] = 0.0 + context_latent, _ = self.context_encoder(context_input, context_mask.float()) + + # Target encoder: full sequence, no gradients; normalized targets. + with torch.no_grad(): + target_latent, _ = self.target_encoder(sequence, event_mask) + target_latent = self.target_norm(target_latent) + + # Predictor input: context positions carry projected context latents; + # target positions carry a location (+scale) aware query. + P = self.predictor_dim + ctx_proj = self.predictor_context_proj(context_latent) # (B, S, P) + positions = torch.arange(S, device=device) + pos_emb = self._sinusoidal_pos_emb(positions, P, device) # (S, P) + pos_emb = pos_emb.unsqueeze(0).to(ctx_proj.dtype) # (1, S, P) + query = self.predictor_query.to(ctx_proj.dtype) + pos_emb # (1, S, P) + if scale_ids is not None and hasattr(self, "scale_embed"): + query = query + self.scale_embed(scale_ids) # (B, S, P) + # Keep the predictor input in the context dtype (under AMP nn.Embedding + # stays fp32, which would otherwise upcast the whole predictor input). + query = query.to(ctx_proj.dtype) + pred_in = torch.where(target_mask.unsqueeze(-1), query, ctx_proj) + + pred_out, _ = self.predictor(pred_in, event_mask) # (B, S, P) + pred_out = self.predictor_norm(pred_out) + pred = self.predictor_out_proj(pred_out) # (B, S, E) + return pred, target_latent + + def forward( + self, + inputs: Optional[dict[str, dict[str, torch.Tensor]]] = None, + target_mask: Optional[torch.Tensor] = None, + context_mask: Optional[torch.Tensor] = None, + feature_keys: Optional[list[str]] = None, + input_processors: Optional[dict[str, Any]] = None, + **raw_kwargs: torch.Tensor | tuple[torch.Tensor, ...], + ) -> dict[str, torch.Tensor]: + if inputs is None: + from .utils import build_unified_inputs_from_batch + fk = feature_keys or getattr(self, "feature_keys", None) + ip = input_processors or getattr(self, "input_processors", None) + if fk is None or ip is None: + raise ValueError( + "When 'inputs' is not provided, both 'feature_keys' and " + "'input_processors' are required (either as arguments or " + "as attributes on the model)." + ) + inputs = build_unified_inputs_from_batch(ip, fk, raw_kwargs) + + emb_out = self.embedding_model(inputs) + sequence = emb_out["sequence"] # (B, S, E) + event_mask = emb_out["mask"] # (B, S) + type_ids = emb_out.get("type_ids") + + if target_mask is None or context_mask is None: + target_mask, context_mask = self._sample_target_blocks(event_mask) + + pred, target_latent = self._encode_and_predict( + sequence, event_mask, target_mask, context_mask + ) + + valid_target = target_mask & event_mask.bool() + if not valid_target.any(): + # Graph-connected zero (keeps every parameter in the autograd graph + # so backward/DDP do not break); contributes no signal this step. + loss = pred.sum() * 0.0 + return { + "loss": loss, + "loss_dict": {"total": 0.0}, + "context_pred": pred[valid_target], + "target_embs": target_latent[valid_target].detach(), + "target_mask": target_mask, + "context_mask": context_mask, + "event_mask": event_mask, + "type_ids": type_ids, + } + + pred_t = pred[valid_target] # (N, E) + tgt_t = target_latent[valid_target].detach() # (N, E) + per_pos_loss = ((pred_t - tgt_t) ** 2).mean(dim=-1) # (N,) + loss = per_pos_loss.mean() + loss_dict = {"total": loss.item()} + if type_ids is not None: + tt = type_ids[valid_target] + for t in tt.unique(): + loss_dict[f"modality_{int(t.item())}"] = per_pos_loss[tt == t].mean().item() + return { + "loss": loss, + "loss_dict": loss_dict, + "context_pred": pred_t, + "target_embs": tgt_t, + "target_mask": target_mask, + "context_mask": context_mask, + "event_mask": event_mask, + "type_ids": type_ids, + } + + +class MultimodalVJEPA(MultimodalIJEPA): + """V-JEPA-style pretrainer for unified multimodal event sequences. + + This extends :class:`MultimodalIJEPA` along the axes that distinguish + V-JEPA (video JEPA) from I-JEPA, adapted to temporal EHR sequences: + + 1. **Multi-scale span (tube) masking.** Target blocks are sampled with + lengths drawn from ``target_block_scales`` (short blocks capture + fine-grained dynamics, long blocks force trend prediction). + 2. **Scale aware predictor.** In addition to the location (positional) + query shared with I-JEPA, each target position adds a learned + ``scale_embed`` for the scale of the block it belongs to. + 3. **Per-position latent prediction.** Inherited from the base class: the + predictor forecasts the EMA-target latent at every masked position. + 4. **Cross-modal target windows (optional).** With + ``require_multimodal_blocks=True`` the sampler prefers contiguous spans + that contain more than one modality, forcing cross-modal temporal + reasoning rather than within-modality interpolation. + + Context and target encoders share architecture (the target encoder is an EMA + copy of the context encoder); an asymmetric design is incompatible with the + EMA-copy update and is therefore intentionally not used. + + Args: + embedding_model: Unified embedding model. + context_encoder: Sequence encoder backbone. + predictor_layers / predictor_heads / predictor_dim: Predictor config. + target_ema_decay / target_ema_end: EMA momentum schedule endpoints. + num_target_blocks: Number of target blocks to sample per sample. + target_block_scales: Candidate block lengths (the multi-scale set). + min_context_len: Minimum visible context positions to keep. + require_multimodal_blocks: Prefer spans spanning >1 modality. + normalize_targets: LayerNorm (no affine) targets before the loss. + """ + + def __init__( + self, + embedding_model: UnifiedMultimodalEmbeddingModel, + context_encoder: nn.Module, + predictor_layers: int = 6, + predictor_heads: int = 8, + predictor_dim: Optional[int] = None, + target_ema_decay: float = 0.996, + target_ema_end: float = 1.0, + num_target_blocks: int = 4, + target_block_scales: tuple[int, ...] = (2, 4, 8), + min_context_len: int = 4, + require_multimodal_blocks: bool = False, + normalize_targets: bool = True, + ): + if not target_block_scales: + raise ValueError("target_block_scales must be non-empty") + mean_len = max(1, int(round(sum(target_block_scales) / len(target_block_scales)))) + super().__init__( + embedding_model=embedding_model, + context_encoder=context_encoder, + predictor_layers=predictor_layers, + predictor_heads=predictor_heads, + predictor_dim=predictor_dim, + target_ema_decay=target_ema_decay, + target_ema_end=target_ema_end, + num_target_blocks=num_target_blocks, + target_block_len=mean_len, + min_context_len=min_context_len, + normalize_targets=normalize_targets, + ) + self.target_block_scales = list(target_block_scales) + self.require_multimodal_blocks = require_multimodal_blocks + + # Scale-aware component of the predictor query (one vector per scale). + self.scale_embed = nn.Embedding(len(self.target_block_scales), self.predictor_dim) + nn.init.trunc_normal_(self.scale_embed.weight, std=0.02) + + def _sample_multiscale_blocks( + self, + event_mask: torch.Tensor, + type_ids: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Sample multi-scale, non-overlapping target spans. + + Returns: + target_mask: ``(B, S)`` bool, True = target position. + context_mask: ``(B, S)`` bool, True = context position. + scale_ids: ``(B, S)`` long, index into ``target_block_scales`` for + each target position (0 elsewhere). + """ + B, S = event_mask.shape + device = event_mask.device + valid = event_mask.bool() + target_mask = torch.zeros_like(event_mask, dtype=torch.bool) + scale_ids = torch.zeros_like(event_mask, dtype=torch.long) + scales = self.target_block_scales + + for b in range(B): + valid_positions = valid[b].nonzero(as_tuple=False).flatten() + n_valid = int(valid_positions.numel()) + if n_valid <= self.min_context_len + 1: + # Too short for span masking: optionally hide one position. + if n_valid > self.min_context_len: + j = int(torch.randint(0, n_valid, (1,), device=device).item()) + target_mask[b, valid_positions[j]] = True + continue + + max_targets = n_valid - self.min_context_len + used = torch.zeros(n_valid, dtype=torch.bool, device=device) + n_blocks = 0 + attempts = 0 + max_attempts = self.num_target_blocks * 20 + while n_blocks < self.num_target_blocks and attempts < max_attempts: + attempts += 1 + s_idx = int(torch.randint(0, len(scales), (1,), device=device).item()) + block_len = min(scales[s_idx], max_targets) + if block_len <= 0: + break + max_start = n_valid - block_len + if max_start < 0: + continue + start = int(torch.randint(0, max_start + 1, (1,), device=device).item()) + span = torch.arange(start, start + block_len, device=device) + if used[span].any(): + continue + if int(used.sum().item()) + block_len > max_targets: + continue + seg_positions = valid_positions[span] + # Prefer cross-modal windows when requested (soft constraint). + if ( + self.require_multimodal_blocks + and type_ids is not None + and attempts < max_attempts // 2 + and int(type_ids[b, seg_positions].unique().numel()) < 2 + ): + continue + used[span] = True + target_mask[b, seg_positions] = True + scale_ids[b, seg_positions] = s_idx + n_blocks += 1 + + context_mask = valid & ~target_mask + return target_mask, context_mask, scale_ids + + def forward( + self, + inputs: Optional[dict[str, dict[str, torch.Tensor]]] = None, + feature_keys: Optional[list[str]] = None, + input_processors: Optional[dict[str, Any]] = None, + **raw_kwargs: torch.Tensor | tuple[torch.Tensor, ...], + ) -> dict[str, torch.Tensor]: + if inputs is None: + from .utils import build_unified_inputs_from_batch + fk = feature_keys or getattr(self, "feature_keys", None) + ip = input_processors or getattr(self, "input_processors", None) + if fk is None or ip is None: + raise ValueError( + "When 'inputs' is not provided, both 'feature_keys' and " + "'input_processors' are required (either as arguments or " + "as attributes on the model)." + ) + inputs = build_unified_inputs_from_batch(ip, fk, raw_kwargs) + + emb_out = self.embedding_model(inputs) + sequence = emb_out["sequence"] # (B, S, E) + event_mask = emb_out["mask"] # (B, S) + type_ids = emb_out.get("type_ids") + + target_mask, context_mask, scale_ids = self._sample_multiscale_blocks( + event_mask, type_ids + ) + + pred, target_latent = self._encode_and_predict( + sequence, event_mask, target_mask, context_mask, scale_ids + ) + + valid_target = target_mask & event_mask.bool() + if not valid_target.any(): + loss = pred.sum() * 0.0 # graph-connected zero (no signal this step) + return { + "loss": loss, + "loss_dict": {"total": 0.0}, + "context_pred": pred[valid_target], + "target_embs": target_latent[valid_target].detach(), + "target_mask": target_mask, + "context_mask": context_mask, + "scale_ids": scale_ids, + "event_mask": event_mask, + "type_ids": type_ids, + } + + pred_t = pred[valid_target] # (N, E) + tgt_t = target_latent[valid_target].detach() # (N, E) + per_pos_loss = ((pred_t - tgt_t) ** 2).mean(dim=-1) + loss = per_pos_loss.mean() + + loss_dict: dict[str, float] = {"total": loss.item()} + scale_t = scale_ids[valid_target] + for s in scale_t.unique(): + block_len = self.target_block_scales[int(s.item())] + loss_dict[f"scale_{block_len}"] = per_pos_loss[scale_t == s].mean().item() + if type_ids is not None: + tt = type_ids[valid_target] + for t in tt.unique(): + loss_dict[f"modality_{int(t.item())}"] = per_pos_loss[tt == t].mean().item() + + return { + "loss": loss, + "loss_dict": loss_dict, + "context_pred": pred_t, + "target_embs": tgt_t, + "target_mask": target_mask, + "context_mask": context_mask, + "scale_ids": scale_ids, + "event_mask": event_mask, + "type_ids": type_ids, + } diff --git a/pyhealth/models/pretrain/mae.py b/pyhealth/models/pretrain/mae.py new file mode 100644 index 000000000..fbdd99af2 --- /dev/null +++ b/pyhealth/models/pretrain/mae.py @@ -0,0 +1,295 @@ +"""Masked Autoencoder (MAE) for unified multimodal event sequences. + +The model embeds heterogeneous temporal features with +:class:`UnifiedMultimodalEmbeddingModel`, masks a fraction of the unified event +sequence, encodes the visible subset with a Transformer backbone, then decodes +the masked positions with a lightweight transformer decoder. The reconstruction +target is the original unified event embedding (time + type + modality token) +before masking. + +References: + Kaiming He et al., "Masked Autoencoders Are Scalable Vision Learners", + CVPR 2022. +""" + +from __future__ import annotations + +import warnings +from typing import Any, Optional + +import torch +import torch.nn as nn + +from ...processors.base_processor import ModalityType +from ..embedding.unified import UnifiedMultimodalEmbeddingModel +from ..transformer import TransformerLayer +from .masking import UnifiedMaskGenerator, apply_mask_token + + +class MultimodalMaskedAutoencoder(nn.Module): + """True MAE pretrainer over a unified temporal event sequence. + + Args: + embedding_model: The unified embedding model that produces a single + temporally-sorted event sequence. + backbone: Transformer (or other sequence encoder) backbone. + decoder_layers: Number of layers in the lightweight reconstruction + decoder. Default 4. + decoder_heads: Attention heads in the decoder. Default 8. + decoder_dim: Hidden dim of the decoder. Defaults to ``embedding_dim``. + mask_ratio: Fraction of valid positions to mask. Default 0.5. + mask_strategy: ``"random"`` or ``"block"``. + per_modality_ratio: Optional override of ``mask_ratio`` per modality + index. + norm_pix_loss: If True, normalize targets by their mean/std before + computing MSE (MAE default for images). + target: Reconstruction target. ``"token"`` (default) reconstructs the + content-only per-event embedding (``token_emb``) before time/type + are added — this is the recommended objective because the time/type + components are largely recoverable from event position and otherwise + dilute the content signal. ``"unified"`` reconstructs the full + composed ``sequence`` (legacy behaviour). + + Inputs: + Same dict expected by ``embedding_model.forward``: + ``{field: {"value": ..., "time": ..., "mask": ...}}``. + + Outputs: + Dict with keys: + - ``loss``: scalar MSE reconstruction loss. + - ``loss_dict``: per-modality MSE breakdown. + - ``pred``: ``(B, S, E)`` decoder predictions for masked positions. + - ``target``: ``(B, S, E)`` reconstruction target. + - ``mask_token``: ``(B, S)`` bool, True = masked. + """ + + def __init__( + self, + embedding_model: UnifiedMultimodalEmbeddingModel, + backbone: nn.Module, + decoder_layers: int = 4, + decoder_heads: int = 8, + decoder_dim: Optional[int] = None, + mask_ratio: float = 0.5, + mask_strategy: str = "random", + per_modality_ratio: Optional[dict[int, float]] = None, + norm_pix_loss: bool = False, + target: str = "token", + ): + super().__init__() + if target not in ("token", "unified"): + raise ValueError(f"target must be 'token' or 'unified', got {target}") + self.target = target + self.embedding_model = embedding_model + self.backbone = backbone + self.embedding_dim = embedding_model.embedding_dim + self.decoder_dim = decoder_dim or self.embedding_dim + + self.mask_generator = UnifiedMaskGenerator( + mask_ratio=mask_ratio, + strategy=mask_strategy, + per_modality_ratio=per_modality_ratio, + ) + self.mask_token = nn.Parameter(torch.zeros(self.embedding_dim)) + nn.init.trunc_normal_(self.mask_token, std=0.02) + + # Transformer decoder: project encoder output to decoder dim, add mask + # tokens, then run a shallow Transformer. + self.encoder_to_decoder = ( + nn.Linear(self.embedding_dim, self.decoder_dim) + if self.decoder_dim != self.embedding_dim + else nn.Identity() + ) + self.decoder_pos_embed = nn.Parameter( + torch.zeros(1, 1, self.decoder_dim) + ) + nn.init.trunc_normal_(self.decoder_pos_embed, std=0.02) + self.decoder = TransformerLayer( + feature_size=self.decoder_dim, + heads=decoder_heads, + dropout=0.0, + num_layers=decoder_layers, + ) + self.decoder_norm = nn.LayerNorm(self.decoder_dim) + self.prediction_head = nn.Linear(self.decoder_dim, self.embedding_dim) + + self.norm_pix_loss = norm_pix_loss + + def forward( + self, + inputs: Optional[dict[str, dict[str, torch.Tensor]]] = None, + mask_token: Optional[torch.Tensor] = None, + feature_keys: Optional[list[str]] = None, + input_processors: Optional[dict[str, Any]] = None, + **raw_kwargs: torch.Tensor | tuple[torch.Tensor, ...], + ) -> dict[str, torch.Tensor]: + if inputs is None: + from .utils import build_unified_inputs_from_batch + fk = feature_keys or getattr(self, "feature_keys", None) + ip = input_processors or getattr(self, "input_processors", None) + if fk is None or ip is None: + raise ValueError( + "When 'inputs' is not provided, both 'feature_keys' and " + "'input_processors' are required (either as arguments or " + "as attributes on the model)." + ) + inputs = build_unified_inputs_from_batch(ip, fk, raw_kwargs) + + # 1. Embed all events. The encoder input is the composed sequence + # (content + time + type); the reconstruction target is content-only by + # default (see ``target``). + emb_out = self.embedding_model(inputs) + sequence = emb_out["sequence"] # (B, S, E) encoder input + target = self._resolve_target(emb_out) # (B, S, E) recon target + event_mask = emb_out["mask"] # (B, S) + type_ids = emb_out.get("type_ids") # (B, S) + B, S, E = sequence.shape + + # 2. Generate masking pattern. + if mask_token is None: + mask_token = self.mask_generator(event_mask, type_ids) + + # 3. Replace masked positions with the learnable mask token and encode. + masked_input = apply_mask_token(sequence, mask_token, self.mask_token) + encoded, _ = self.backbone(masked_input, event_mask) # (B, S, E) + encoded = self.encoder_to_decoder(encoded) # (B, S, D) + + # 4. Decoder: process full sequence. Per-event localization already + # comes from the unified embedding's sinusoidal time embedding; this is + # just a small global learnable decoder bias. + decoder_input = encoded + self.decoder_pos_embed + decoded, _ = self.decoder(decoder_input, event_mask) # (B, S, D) + decoded = self.decoder_norm(decoded) + pred = self.prediction_head(decoded) # (B, S, E) + + # 5. Reconstruction loss on masked positions only. + loss, loss_dict = self._reconstruction_loss( + pred, target, mask_token, event_mask, type_ids + ) + + return { + "loss": loss, + "loss_dict": loss_dict, + "pred": pred, + "target": target, + "mask_token": mask_token, + "event_mask": event_mask, + "type_ids": type_ids, + } + + def _resolve_target(self, emb_out: dict[str, torch.Tensor]) -> torch.Tensor: + """Pick the reconstruction target (content-only by default). + + The target is always **detached**: it is a function of the (trainable) + embedding model, so back-propagating through it lets the model trivially + shrink the target to zero (representation collapse). ``token`` targets + are additionally normalized in the loss to be scale-invariant. + """ + if self.target == "token": + if "token_emb" in emb_out: + return emb_out["token_emb"].detach() + warnings.warn( + "target='token' requested but the embedding model did not return " + "'token_emb'; falling back to the composed 'sequence' target.", + stacklevel=2, + ) + return emb_out["sequence"].detach() + + def _reconstruction_loss( + self, + pred: torch.Tensor, + target: torch.Tensor, + mask_token: torch.Tensor, + event_mask: torch.Tensor, + type_ids: Optional[torch.Tensor], + ) -> tuple[torch.Tensor, dict[str, float]]: + """Compute per-modality normalized MSE on masked valid positions.""" + valid_mask = event_mask.bool() & mask_token + if not valid_mask.any(): + # Graph-connected zero so every parameter still receives a (zero) + # gradient; a disconnected leaf tensor breaks DDP / AMP GradScaler. + return pred.sum() * 0.0, {"total": 0.0} + + pred_masked = pred[valid_mask] # (N, E) + target_masked = target[valid_mask] # (N, E) + + # Normalize when requested, and always for the learnable ``token`` + # target — a scale-invariant target removes the remaining slow, + # weight-decay-driven shrink path toward collapse. + if self.norm_pix_loss or self.target == "token": + mean = target_masked.mean(dim=-1, keepdim=True) + var = target_masked.var(dim=-1, keepdim=True, unbiased=False) + target_masked = (target_masked - mean) / (var + 1e-6).sqrt() + + per_pos_loss = ((pred_masked - target_masked) ** 2).mean(dim=-1) # (N,) + + loss_dict: dict[str, float] = {} + if type_ids is not None: + type_ids_masked = type_ids[valid_mask] + unique_types = type_ids_masked.unique() + for t in unique_types: + name = f"modality_{int(t.item())}" + loss_dict[name] = per_pos_loss[type_ids_masked == t].mean().item() + + loss = per_pos_loss.mean() + loss_dict["total"] = loss.item() + return loss, loss_dict + + +class PerModalityMAEDecoder(nn.Module): + """Optional add-on that decodes unified embeddings back to raw modality features. + + This is kept separate from :class:`MultimodalMaskedAutoencoder` so that the + default MAE can predict in the unified embedding space (stable, fast), + while experiments that want true raw-value reconstruction can attach + per-modality heads without touching the core MAE code. + + Args: + embedding_dim: Dimension of unified event embeddings. + output_specs: Dict mapping modality index to output dimension and + prediction type. Example:: + + {0: ("numeric", 10), 1: ("code", vocab_size)} + + hidden_dim: Hidden size of the small MLP decoder per modality. + """ + + def __init__( + self, + embedding_dim: int, + output_specs: dict[int, tuple[str, int]], + hidden_dim: int = 256, + ): + super().__init__() + self.heads = nn.ModuleDict() + for mod_idx, (task, out_dim) in output_specs.items(): + key = str(mod_idx) + if task == "numeric": + self.heads[key] = nn.Sequential( + nn.Linear(embedding_dim, hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, out_dim), + ) + elif task == "code": + self.heads[key] = nn.Sequential( + nn.Linear(embedding_dim, hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, out_dim), + ) + else: + raise ValueError(f"Unknown per-modality task: {task}") + + def forward( + self, + unified_embedding: torch.Tensor, + type_ids: torch.Tensor, + ) -> dict[int, torch.Tensor]: + """Return per-modality predictions keyed by modality index.""" + outputs: dict[int, torch.Tensor] = {} + for key, head in self.heads.items(): + mod_idx = int(key) + mask = type_ids == mod_idx + if not mask.any(): + continue + outputs[mod_idx] = head(unified_embedding[mask]) + return outputs diff --git a/pyhealth/models/pretrain/masking.py b/pyhealth/models/pretrain/masking.py new file mode 100644 index 000000000..c45600f37 --- /dev/null +++ b/pyhealth/models/pretrain/masking.py @@ -0,0 +1,183 @@ +"""Masking utilities for self-supervised pretraining on unified event sequences. + +A unified event sequence has shape ``(B, S_total, E)`` with an accompanying +validity mask ``(B, S_total)`` and modality type ids ``(B, S_total)``. The +collators below generate boolean ``mask_token`` tensors that select which +positions are hidden during pretraining. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + + +class UnifiedMaskGenerator: + """Generate masking patterns for a unified temporal event sequence. + + Args: + mask_ratio: Fraction of valid (non-padding) positions to mask. + strategy: ``"random"`` or ``"block"``. Block masking hides contiguous + spans, which is closer to MAE/SimMIM and more realistic for EHR + (a missing lab window rather than random single events). + min_block_len: Minimum span length for block masking. + max_block_len: Maximum span length for block masking. + per_modality_ratio: Optional dict ``{modality_index: ratio}`` that + overrides ``mask_ratio`` for specific modality types. Useful when + text should be masked less aggressively than labs. + seed: Not used; callers should set the global/random generator for + reproducibility. + + Shape: + Input ``mask`` is ``(B, S)`` with 1 = valid, 0 = padding. + Output ``mask_token`` is ``(B, S)`` bool, True = hide this position. + """ + + def __init__( + self, + mask_ratio: float = 0.5, + strategy: str = "random", + min_block_len: int = 3, + max_block_len: int = 12, + per_modality_ratio: Optional[dict[int, float]] = None, + ): + if not 0.0 <= mask_ratio < 1.0: + raise ValueError(f"mask_ratio must be in [0, 1), got {mask_ratio}") + self.mask_ratio = mask_ratio + self.strategy = strategy + self.min_block_len = min_block_len + self.max_block_len = max_block_len + self.per_modality_ratio = per_modality_ratio or {} + + def __call__( + self, + mask: torch.Tensor, + type_ids: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Return a boolean mask_token tensor of the same shape as ``mask``.""" + if self.strategy == "random": + return self._random_mask(mask, type_ids) + if self.strategy == "block": + return self._block_mask(mask, type_ids) + raise ValueError(f"Unknown masking strategy: {self.strategy}") + + def _effective_ratio( + self, + ref: torch.Tensor, + type_ids: Optional[torch.Tensor], + default_ratio: float, + ) -> torch.Tensor: + """Build a per-position target mask ratio with shape ``ref.shape``.""" + if type_ids is None or not self.per_modality_ratio: + return torch.full_like(ref, default_ratio, dtype=torch.float32) + ratio = torch.full_like(type_ids, default_ratio, dtype=torch.float32) + for mod_idx, mod_ratio in self.per_modality_ratio.items(): + ratio = torch.where(type_ids == mod_idx, mod_ratio, ratio) + return ratio + + def _random_mask( + self, + mask: torch.Tensor, + type_ids: Optional[torch.Tensor], + ) -> torch.Tensor: + """Independent Bernoulli masking over valid positions.""" + valid = mask.bool() + ratio = self._effective_ratio(mask, type_ids, self.mask_ratio) + probs = torch.zeros_like(mask, dtype=torch.float32) + probs[valid] = ratio[valid] + mask_token = torch.bernoulli(probs).bool() & valid + # Floor: every sample that has valid positions must mask at least one, + # otherwise that sample contributes no reconstruction/prediction signal. + empty_rows = valid.any(dim=1) & ~mask_token.any(dim=1) + for b in empty_rows.nonzero(as_tuple=False).flatten(): + vp = valid[b].nonzero(as_tuple=False).flatten() + j = vp[torch.randint(0, vp.numel(), (1,), device=mask.device)] + mask_token[b, j] = True + return mask_token + + def _block_mask( + self, + mask: torch.Tensor, + type_ids: Optional[torch.Tensor], + ) -> torch.Tensor: + """Sample contiguous spans until the target ratio is reached per sample.""" + B, S = mask.shape + device = mask.device + mask_token = torch.zeros_like(mask, dtype=torch.bool) + + for b in range(B): + valid_positions = mask[b].nonzero(as_tuple=False).flatten() + if valid_positions.numel() == 0: + continue + + # Determine effective ratio per position. + if type_ids is not None and self.per_modality_ratio: + ratios = self._effective_ratio( + mask[b : b + 1], type_ids[b : b + 1], self.mask_ratio + ).squeeze(0) + target_count = int( + round(ratios[valid_positions].float().mean().item() * valid_positions.numel()) + ) + else: + target_count = int(round(self.mask_ratio * valid_positions.numel())) + + chosen = torch.zeros_like(valid_positions, dtype=torch.bool) + attempts = 0 + max_attempts = target_count * 4 + 10 + while chosen.sum().item() < target_count and attempts < max_attempts: + attempts += 1 + remaining = target_count - int(chosen.sum().item()) + block_len = int( + torch.randint( + self.min_block_len, + self.max_block_len + 1, + (1,), + device=device, + ).item() + ) + # Cap the block to the remaining budget so we never overshoot and + # never need a random trim (random trimming shatters the spans, + # defeating the purpose of block masking). + block_len = min(block_len, remaining) + start_idx = int( + torch.randint( + 0, max(1, valid_positions.numel() - block_len + 1), (1,), device=device + ).item() + ) + chosen[start_idx : start_idx + block_len] = True + + # `chosen` marks contiguous runs in valid-event order; map back to the + # actual sequence positions (kept contiguous over valid events). + mask_token[b, valid_positions[chosen]] = True + + return mask_token + + +def apply_mask_token( + sequence: torch.Tensor, + mask_token: torch.Tensor, + learnable_mask_token: torch.Tensor, +) -> torch.Tensor: + """Replace masked positions in ``sequence`` with a learnable mask token. + + Args: + sequence: ``(B, S, E)`` unified event embeddings. + mask_token: ``(B, S)`` bool, True = hide. + learnable_mask_token: ``(E,)`` shared mask embedding. + + Returns: + ``(B, S, E)`` sequence with masked positions replaced. + """ + masked = sequence.clone() + masked[mask_token] = learnable_mask_token + return masked + + +def random_mask_like( + mask: torch.Tensor, + mask_ratio: float = 0.5, +) -> torch.Tensor: + """Convenience one-liner for independent random masking.""" + gen = UnifiedMaskGenerator(mask_ratio=mask_ratio, strategy="random") + return gen(mask) diff --git a/pyhealth/models/pretrain/rope.py b/pyhealth/models/pretrain/rope.py new file mode 100644 index 000000000..1beb52b82 --- /dev/null +++ b/pyhealth/models/pretrain/rope.py @@ -0,0 +1,266 @@ +"""Rotary Position Embedding (RoPE) utilities. + +RoPE encodes relative position by rotating query/key vectors in 2D subspaces. +It is especially useful for long clinical sequences because it generalizes to +lengths longer than those seen during training and supports extrapolation +techniques such as NTK-aware scaling or YaRN. + +References: + Jianlin Su et al., "RoFormer: Enhanced Transformer with Rotary Position + Embedding", Neurocomputing 2024. +""" + +from __future__ import annotations + +import math +from typing import Optional + +import torch +import torch.nn as nn + + +class RotaryPositionEmbedding(nn.Module): + """Rotary position embedding for sequences. + + Args: + dim: Head dimension (must be even). + max_seq_len: Maximum sequence length for which to precompute angles. + base: Base for the inverse frequency computation. Default 10000. + scaling_factor: Multiplicative scaling for sequence-length + extrapolation (e.g., NTK-aware scaling). Default 1.0. + + Shape: + Input: ``(..., seq_len, head_dim)`` + Output: ``(..., seq_len, head_dim)`` rotated by position. + """ + + def __init__( + self, + dim: int, + max_seq_len: int = 8192, + base: float = 10000.0, + scaling_factor: float = 1.0, + ): + super().__init__() + if dim % 2 != 0: + raise ValueError(f"RoPE dim must be even, got {dim}") + self.dim = dim + self.max_seq_len = max_seq_len + self.base = base + self.scaling_factor = scaling_factor + + inv_freq = self._compute_inv_freq() + self.register_buffer("inv_freq", inv_freq, persistent=False) + + # Precompute cos/sin caches. + self._update_cos_sin_cache(max_seq_len, device=inv_freq.device) + + def _compute_inv_freq(self) -> torch.Tensor: + # Standard RoPE: inv_freq_i = 1 / base^(2i/dim). NTK-aware extrapolation + # rescales the base by scaling_factor^(dim/(dim-2)) so that longer + # sequences interpolate smoothly (scaling_factor=1.0 -> vanilla RoPE). + base = self.base + if self.scaling_factor != 1.0: + base = base * (self.scaling_factor ** (self.dim / (self.dim - 2))) + exponent = torch.arange(0, self.dim, 2).float() / self.dim + return 1.0 / (base ** exponent) + + def _update_cos_sin_cache(self, seq_len: int, device: torch.device) -> None: + if ( + hasattr(self, "cos_cached") + and self.cos_cached.shape[1] >= seq_len # axis 1 is the cached seq len + and self.cos_cached.device == device + ): + return + positions = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) + freqs = torch.outer(positions, self.inv_freq) # (seq_len, dim/2) + emb = torch.cat([freqs, freqs], dim=-1) # (seq_len, dim) + self.register_buffer("cos_cached", emb.cos()[None, :, :], persistent=False) + self.register_buffer("sin_cached", emb.sin()[None, :, :], persistent=False) + + def _rotate_half(self, x: torch.Tensor) -> torch.Tensor: + """Rotate the last dimension by swapping pairs and negating.""" + x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] + return torch.cat([-x2, x1], dim=-1) + + def forward(self, x: torch.Tensor, seq_len: Optional[int] = None) -> torch.Tensor: + """Apply rotary embedding to ``x`` of shape ``(..., seq_len, dim)``.""" + if seq_len is None: + seq_len = x.shape[-2] + self._update_cos_sin_cache(seq_len, device=x.device) + cos = self.cos_cached[:, :seq_len, :] + sin = self.sin_cached[:, :seq_len, :] + return x * cos + self._rotate_half(x) * sin + + +class RoPEMultiHeadedAttention(nn.Module): + """Multi-head attention with RoPE applied to Q and K. + + This is a drop-in replacement for + :class:`pyhealth.models.transformer.MultiHeadedAttention`. + + Args: + h: Number of attention heads. + d_model: Model dimensionality (must be divisible by ``h``). + dropout: Dropout probability on attention weights. + rope_max_seq_len: Max sequence length for RoPE cache. + rope_base: RoPE inverse-frequency base. + rope_scaling: RoPE scaling factor for extrapolation. + """ + + def __init__( + self, + h: int, + d_model: int, + dropout: float = 0.1, + rope_max_seq_len: int = 8192, + rope_base: float = 10000.0, + rope_scaling: float = 1.0, + ): + super().__init__() + if d_model % h != 0: + raise ValueError("d_model must be divisible by h") + self.d_k = d_model // h + self.h = h + self.linear_layers = nn.ModuleList( + [nn.Linear(d_model, d_model, bias=False) for _ in range(3)] + ) + self.output_linear = nn.Linear(d_model, d_model, bias=False) + self.dropout = nn.Dropout(p=dropout) + self.rope = RotaryPositionEmbedding( + dim=self.d_k, + max_seq_len=rope_max_seq_len, + base=rope_base, + scaling_factor=rope_scaling, + ) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + mask: Optional[torch.Tensor] = None, + register_hook: bool = False, + ) -> torch.Tensor: + batch_size = query.size(0) + query, key, value = [ + l(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2) + for l, x in zip(self.linear_layers, (query, key, value)) + ] + + # Apply RoPE to Q and K; V is left unchanged. + query = self.rope(query) + key = self.rope(key) + + if mask is not None: + # (B, S, S) -> (B, 1, S, S) so it broadcasts across heads (mirrors + # MultiHeadedAttention). Without this, multi-head crashes and the + # B == h case silently mis-masks. + mask = mask.unsqueeze(1) + + scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.d_k) + if mask is not None: + # dtype min keeps this safe under fp16 AMP (-1e9 overflows half). + scores = scores.masked_fill(mask == 0, torch.finfo(scores.dtype).min) + p_attn = torch.softmax(scores, dim=-1) + if mask is not None: + p_attn = p_attn.masked_fill(mask == 0, 0) + p_attn = self.dropout(p_attn) + + x = torch.matmul(p_attn, value) + x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k) + return self.output_linear(x) + + +class RoPETransformerBlock(nn.Module): + """Transformer block using :class:`RoPEMultiHeadedAttention`.""" + + def __init__( + self, + hidden: int, + attn_heads: int, + dropout: float, + rope_max_seq_len: int = 8192, + rope_base: float = 10000.0, + rope_scaling: float = 1.0, + ): + super().__init__() + from pyhealth.models.transformer import PositionwiseFeedForward, SublayerConnection + + self.attention = RoPEMultiHeadedAttention( + h=attn_heads, + d_model=hidden, + dropout=dropout, + rope_max_seq_len=rope_max_seq_len, + rope_base=rope_base, + rope_scaling=rope_scaling, + ) + self.feed_forward = PositionwiseFeedForward( + d_model=hidden, d_ff=4 * hidden, dropout=dropout + ) + self.input_sublayer = SublayerConnection(size=hidden, dropout=dropout) + self.output_sublayer = SublayerConnection(size=hidden, dropout=dropout) + self.dropout = nn.Dropout(p=dropout) + + def forward( + self, + x: torch.Tensor, + mask: Optional[torch.Tensor] = None, + register_hook: bool = False, + ) -> torch.Tensor: + x = self.input_sublayer( + x, lambda _x: self.attention(_x, _x, _x, mask=mask, register_hook=register_hook) + ) + x = self.output_sublayer(x, lambda _x: self.feed_forward(_x, mask=mask)) + return self.dropout(x) + + +class RoPETransformerLayer(nn.Module): + """RoPE-enabled Transformer layer matching the interface of + :class:`pyhealth.models.transformer.TransformerLayer`. + + When used as the SSL backbone, the unified embedding model can optionally + omit its sinusoidal time embedding (set ``time_embedding="none"`` once + supported) because RoPE encodes relative position directly in attention. + """ + + def __init__( + self, + feature_size: int, + heads: int = 1, + dropout: float = 0.5, + num_layers: int = 1, + rope_max_seq_len: int = 8192, + rope_base: float = 10000.0, + rope_scaling: float = 1.0, + ): + super().__init__() + self.feature_size = feature_size + self.heads = heads + self.dropout = dropout + self.num_layers = num_layers + self.transformer = nn.ModuleList( + [ + RoPETransformerBlock( + hidden=feature_size, + attn_heads=heads, + dropout=dropout, + rope_max_seq_len=rope_max_seq_len, + rope_base=rope_base, + rope_scaling=rope_scaling, + ) + for _ in range(num_layers) + ] + ) + + def forward( + self, + x: torch.Tensor, + mask: Optional[torch.Tensor] = None, + register_hook: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + if mask is not None: + mask = torch.einsum("ab,ac->abc", mask, mask) + for transformer in self.transformer: + x = transformer(x, mask, register_hook) + return x, x[:, 0, :] diff --git a/pyhealth/models/pretrain/simmim.py b/pyhealth/models/pretrain/simmim.py new file mode 100644 index 000000000..71e9a2d9a --- /dev/null +++ b/pyhealth/models/pretrain/simmim.py @@ -0,0 +1,181 @@ +"""SimMIM-style pretraining for unified multimodal event sequences. + +Unlike MAE, SimMIM feeds the *full* sequence (with a learnable mask token at +masked positions) through the encoder and applies a simple linear head on the +final feature map to reconstruct the original per-modality token embedding. +This avoids a separate decoder transformer and is therefore cheaper per step. + +References: + Zhenda Xie et al., "SimMIM: A Simple Framework for Masked Image Modeling", + CVPR 2022. +""" + +from __future__ import annotations + +import warnings +from typing import Any, Optional + +import torch +import torch.nn as nn + +from ..embedding.unified import UnifiedMultimodalEmbeddingModel +from .masking import UnifiedMaskGenerator, apply_mask_token + + +class MultimodalSimMIM(nn.Module): + """SimMIM pretrainer over a unified temporal event sequence. + + Args: + embedding_model: Unified embedding model. + backbone: Sequence encoder (e.g., TransformerLayer, Mamba stack). + mask_ratio: Fraction of valid positions to mask. + mask_strategy: ``"random"`` or ``"block"``. + per_modality_ratio: Optional per-modality mask ratio overrides. + target: What to reconstruct. ``"token"`` (default) predicts the + content-only per-event embedding (``token_emb``) *before* time/type + are added — the recommended objective, since time/type are largely + recoverable from position and otherwise dilute the content signal. + ``"unified"`` predicts the full composed embedding (legacy). + norm_targets: Normalize targets by mean/std before MSE. + + Inputs: + Same dict as ``UnifiedMultimodalEmbeddingModel.forward``. + + Outputs: + Dict with keys ``loss``, ``loss_dict``, ``pred``, ``target``, + ``mask_token``, ``event_mask``, ``type_ids``. + """ + + def __init__( + self, + embedding_model: UnifiedMultimodalEmbeddingModel, + backbone: nn.Module, + mask_ratio: float = 0.5, + mask_strategy: str = "random", + per_modality_ratio: Optional[dict[int, float]] = None, + target: str = "token", + norm_targets: bool = False, + ): + super().__init__() + if target not in ("token", "unified"): + raise ValueError(f"target must be 'token' or 'unified', got {target}") + self.embedding_model = embedding_model + self.backbone = backbone + self.embedding_dim = embedding_model.embedding_dim + self.target = target + self.norm_targets = norm_targets + + self.mask_generator = UnifiedMaskGenerator( + mask_ratio=mask_ratio, + strategy=mask_strategy, + per_modality_ratio=per_modality_ratio, + ) + self.mask_token = nn.Parameter(torch.zeros(self.embedding_dim)) + nn.init.trunc_normal_(self.mask_token, std=0.02) + + # SimMIM uses a single linear prediction head. + self.head = nn.Linear(self.embedding_dim, self.embedding_dim) + + def forward( + self, + inputs: Optional[dict[str, dict[str, torch.Tensor]]] = None, + mask_token: Optional[torch.Tensor] = None, + feature_keys: Optional[list[str]] = None, + input_processors: Optional[dict[str, Any]] = None, + **raw_kwargs: torch.Tensor | tuple[torch.Tensor, ...], + ) -> dict[str, torch.Tensor]: + if inputs is None: + from .utils import build_unified_inputs_from_batch + fk = feature_keys or getattr(self, "feature_keys", None) + ip = input_processors or getattr(self, "input_processors", None) + if fk is None or ip is None: + raise ValueError( + "When 'inputs' is not provided, both 'feature_keys' and " + "'input_processors' are required (either as arguments or " + "as attributes on the model)." + ) + inputs = build_unified_inputs_from_batch(ip, fk, raw_kwargs) + + emb_out = self.embedding_model(inputs) + sequence = emb_out["sequence"] # (B, S, E) encoder input + target = self._resolve_target(emb_out) # (B, S, E) recon target + event_mask = emb_out["mask"] # (B, S) + type_ids = emb_out.get("type_ids") # (B, S) + + if mask_token is None: + mask_token = self.mask_generator(event_mask, type_ids) + + # SimMIM: feed full sequence with mask tokens. + masked_input = apply_mask_token(sequence, mask_token, self.mask_token) + encoded, _ = self.backbone(masked_input, event_mask) # (B, S, E) + pred = self.head(encoded) # (B, S, E) + + loss, loss_dict = self._reconstruction_loss( + pred, target, mask_token, event_mask, type_ids + ) + + return { + "loss": loss, + "loss_dict": loss_dict, + "pred": pred, + "target": target, + "mask_token": mask_token, + "event_mask": event_mask, + "type_ids": type_ids, + } + + def _resolve_target(self, emb_out: dict[str, torch.Tensor]) -> torch.Tensor: + """Pick the reconstruction target (content-only by default). + + Always **detached**: the target is a function of the trainable embedding + model, so back-propagating through it lets the model shrink the target to + zero (representation collapse). ``token`` targets are also normalized in + the loss for scale-invariance. + """ + if self.target == "token": + if "token_emb" in emb_out: + return emb_out["token_emb"].detach() + warnings.warn( + "target='token' requested but the embedding model did not return " + "'token_emb'; falling back to the composed 'sequence' target.", + stacklevel=2, + ) + return emb_out["sequence"].detach() + + def _reconstruction_loss( + self, + pred: torch.Tensor, + target: torch.Tensor, + mask_token: torch.Tensor, + event_mask: torch.Tensor, + type_ids: Optional[torch.Tensor], + ) -> tuple[torch.Tensor, dict[str, float]]: + valid_mask = event_mask.bool() & mask_token + if not valid_mask.any(): + # Graph-connected zero so every parameter still receives a (zero) + # gradient; a disconnected leaf tensor breaks DDP / AMP GradScaler. + return pred.sum() * 0.0, {"total": 0.0} + + pred_masked = pred[valid_mask] + target_masked = target[valid_mask] + + # Normalize when requested, and always for the learnable ``token`` + # target (scale-invariance removes the slow shrink-to-collapse path). + if self.norm_targets or self.target == "token": + mean = target_masked.mean(dim=-1, keepdim=True) + var = target_masked.var(dim=-1, keepdim=True, unbiased=False) + target_masked = (target_masked - mean) / (var + 1e-6).sqrt() + + per_pos_loss = ((pred_masked - target_masked) ** 2).mean(dim=-1) + + loss_dict: dict[str, float] = {} + if type_ids is not None: + type_ids_masked = type_ids[valid_mask] + unique_types = type_ids_masked.unique() + for t in unique_types: + name = f"modality_{int(t.item())}" + loss_dict[name] = per_pos_loss[type_ids_masked == t].mean().item() + + loss = per_pos_loss.mean() + loss_dict["total"] = loss.item() + return loss, loss_dict diff --git a/pyhealth/models/pretrain/trainer.py b/pyhealth/models/pretrain/trainer.py new file mode 100644 index 000000000..44078e37a --- /dev/null +++ b/pyhealth/models/pretrain/trainer.py @@ -0,0 +1,543 @@ +"""Self-supervised pretraining trainer. + +A lightweight trainer specialized for reconstruction / latent-prediction +objectives. It is intentionally separate from :class:`pyhealth.trainer.Trainer` +because SSL has no labels, validation metrics, or classification head. + +Supports single-GPU and multi-GPU (DDP) training. For DDP, launch with:: + + torchrun --nproc_per_node=4 scripts/pretrain_ssl.py --config ... +""" + +from __future__ import annotations + +import json +import logging +import math +import os +import time +from pathlib import Path +from typing import Callable, Dict, List, Optional, Type + +import torch +import torch.distributed as dist +from torch import nn +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import Optimizer +from torch.utils.data import DataLoader, DistributedSampler, IterableDataset +from tqdm import tqdm +from tqdm.autonotebook import trange + +from pyhealth import _wandb +from pyhealth.utils import create_directory + +logger = logging.getLogger(__name__) + + +def set_logger(log_path: str) -> None: + create_directory(log_path) + log_filename = os.path.join(log_path, "log.txt") + handler = logging.FileHandler(log_filename) + formatter = logging.Formatter("%(asctime)s %(message)s", "%Y-%m-%d %H:%M:%S") + handler.setFormatter(formatter) + logger.addHandler(handler) + + +def _vram_stats(device: str) -> Dict[str, float]: + if not torch.cuda.is_available() or not str(device).startswith("cuda"): + return {} + allocated = torch.cuda.memory_allocated(device) / 1024**2 + peak = torch.cuda.max_memory_allocated(device) / 1024**2 + return {"vram_allocated_mb": allocated, "vram_peak_mb": peak} + + +def _is_ddp() -> bool: + return ( + "RANK" in os.environ + and "WORLD_SIZE" in os.environ + and int(os.environ["WORLD_SIZE"]) > 1 + ) + + +def _get_rank() -> int: + return int(os.environ.get("RANK", 0)) + + +def _get_local_rank() -> int: + return int(os.environ.get("LOCAL_RANK", 0)) + + +def _is_main_process() -> bool: + return _get_rank() == 0 + + +class PretrainTrainer: + """Trainer for self-supervised pretraining objectives. + + Args: + model: Pretraining model (MAE, SimMIM, I-JEPA, ...). + device: Device to use. Auto-detected if None. Under DDP this is set + automatically to the local GPU. + enable_logging: Whether to write ``log.txt`` and ``metrics_history.json``. + Only the main process writes logs. + output_path: Root directory for checkpoints and logs. + exp_name: Experiment subdirectory name. + ema_update_fn: Optional callable invoked once per training step to + update an EMA target network (used by I-JEPA). + ema_update_every: Call ``ema_update_fn`` every N steps. Default 1. + use_ddp: If True, use DistributedDataParallel. Defaults to True when + launched via ``torchrun`` / ``RANK`` env vars. + """ + + def __init__( + self, + model: nn.Module, + device: Optional[str] = None, + enable_logging: bool = True, + output_path: Optional[str] = None, + exp_name: Optional[str] = None, + ema_update_fn: Optional[Callable[[], None]] = None, + ema_update_every: int = 1, + use_ddp: Optional[bool] = None, + ): + self._ddp = use_ddp if use_ddp is not None else _is_ddp() + if self._ddp: + if not dist.is_initialized(): + dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo") + self.rank = _get_rank() + self.world_size = int(os.environ["WORLD_SIZE"]) + device = f"cuda:{_get_local_rank()}" if torch.cuda.is_available() else "cpu" + else: + self.rank = 0 + self.world_size = 1 + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + + self.model = model + self.device = device + self.ema_update_fn = ema_update_fn + self.ema_update_every = ema_update_every + + if enable_logging and _is_main_process(): + if output_path is None: + output_path = os.path.join(os.getcwd(), "output") + if exp_name is None: + exp_name = time.strftime("%Y%m%d-%H%M%S") + self.exp_path = os.path.join(output_path, exp_name) + set_logger(self.exp_path) + else: + self.exp_path = None + + self.model.to(self.device) + if self._ddp: + self.model = DDP( + self.model, + device_ids=[_get_local_rank()] if torch.cuda.is_available() else None, + output_device=_get_local_rank() if torch.cuda.is_available() else None, + find_unused_parameters=False, + ) + + if _is_main_process(): + logger.info(self.model) + logger.info(f"Device: {self.device}") + if self._ddp: + logger.info(f"DDP world size: {self.world_size}") + + def train( + self, + train_dataloader: DataLoader, + epochs: int = 10, + optimizer_class: Type[Optimizer] = torch.optim.AdamW, + optimizer_params: Optional[Dict[str, object]] = None, + steps_per_epoch: Optional[int] = None, + weight_decay: float = 0.05, + max_grad_norm: Optional[float] = 1.0, + scheduler: Optional[str] = None, + warmup_steps: int = 0, + save_every_n_epochs: int = 1, + grad_accumulation_steps: int = 1, + use_amp: bool = False, + val_dataloader: Optional[DataLoader] = None, + epoch_callback: Optional[Callable[[int, Dict], None]] = None, + ) -> List[Dict[str, object]]: + """Run SSL pretraining. + + Args: + train_dataloader: Dataloader yielding batches compatible with the + pretraining model's ``forward``. Under DDP this should be + paired with a ``DistributedSampler`` by the caller; if a plain + sampler is detected, the trainer wraps it automatically. + epochs: Number of epochs. + optimizer_class: Optimizer class. + optimizer_params: Optimizer kwargs. Defaults to ``{"lr": 1e-4}``. + steps_per_epoch: If None, uses ``len(train_dataloader)``. + weight_decay: Weight decay. Applied via param-group split. + max_grad_norm: Gradient clipping. None disables. + scheduler: ``"cosine"`` or None. + warmup_steps: Linear warmup steps. + save_every_n_epochs: Save a checkpoint every N epochs. + grad_accumulation_steps: Number of forward/backward steps to + accumulate before an optimizer step. Effective batch size is + ``batch_size * world_size * grad_accumulation_steps``. + use_amp: Use automatic mixed precision (torch.cuda.amp) on CUDA. + + Returns: + List of per-epoch metric dicts (only from the main process when DDP). + """ + if optimizer_params is None: + optimizer_params = {"lr": 1e-4} + + no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"] + + def _decayed(n): + return not any(nd in n for nd in no_decay) + + # Unwrap DDP module for parameter grouping if needed. + raw_model = self.model.module if isinstance(self.model, DDP) else self.model + param_groups = [ + {"params": [p for n, p in raw_model.named_parameters() if _decayed(n)], + "weight_decay": weight_decay}, + {"params": [p for n, p in raw_model.named_parameters() if not _decayed(n)], + "weight_decay": 0.0}, + ] + optimizer = optimizer_class(param_groups, **optimizer_params) + + # Auto-wrap sampler if not already distributed. IterableDataset (e.g. + # litdata StreamingDataset) shards across DDP ranks internally and cannot + # take a sampler, so skip the wrap there. + sampler = train_dataloader.sampler + _iterable = isinstance(train_dataloader.dataset, IterableDataset) + if self._ddp and not _iterable and not isinstance(sampler, DistributedSampler): + train_dataloader = DataLoader( + train_dataloader.dataset, + batch_size=train_dataloader.batch_size, + sampler=DistributedSampler( + train_dataloader.dataset, + num_replicas=self.world_size, + rank=self.rank, + shuffle=True, + ), + num_workers=train_dataloader.num_workers, + collate_fn=train_dataloader.collate_fn, + pin_memory=train_dataloader.pin_memory, + drop_last=train_dataloader.drop_last, + ) + + total_steps = epochs * (steps_per_epoch or len(train_dataloader)) + # The scheduler steps once per *optimizer* step, not per micro-step, so + # the horizon must be divided by the grad-accumulation factor. + total_optim_steps = max(1, total_steps // max(1, grad_accumulation_steps)) + + if scheduler == "cosine" or warmup_steps > 0: + warmup = max(0, int(warmup_steps)) + + def _lr_lambda(opt_step: int) -> float: + # Linear warmup, then (optionally) cosine decay to ~0. + if warmup > 0 and opt_step < warmup: + return (opt_step + 1) / warmup + if scheduler == "cosine": + progress = (opt_step - warmup) / max(1, total_optim_steps - warmup) + return 0.5 * (1.0 + math.cos(math.pi * min(1.0, progress))) + return 1.0 + + sched = torch.optim.lr_scheduler.LambdaLR(optimizer, _lr_lambda) + else: + sched = None + + use_amp = use_amp and torch.cuda.is_available() + scaler = torch.amp.GradScaler("cuda", enabled=use_amp) if use_amp else None + + if _is_main_process(): + logger.info("SSL pretraining:") + logger.info(f"Batch size: {train_dataloader.batch_size}") + logger.info(f"Optimizer: {optimizer_class}") + logger.info(f"Optimizer params: {optimizer_params}") + logger.info(f"Weight decay: {weight_decay}") + logger.info(f"Max grad norm: {max_grad_norm}") + logger.info(f"Epochs: {epochs}") + logger.info(f"Grad accumulation steps: {grad_accumulation_steps}") + logger.info(f"AMP: {use_amp}") + logger.info(f"EMA update every: {self.ema_update_every} steps") + + data_iterator = iter(train_dataloader) + if steps_per_epoch is None: + steps_per_epoch = len(train_dataloader) + global_step = 0 # counts micro (forward/backward) steps + optimizer_step = 0 # counts actual optimizer updates + metrics_history: List[Dict[str, object]] = [] + train_start = time.perf_counter() + + # Resume from a prior (e.g. preempted/requeued) run: restore model + + # optimizer/scheduler state + epoch so it continues instead of restarting. + start_epoch = 0 + if self.exp_path is not None: + _ck = os.path.join(self.exp_path, "last.ckpt") + _rs = os.path.join(self.exp_path, "_resume.pt") + if os.path.isfile(_ck) and os.path.isfile(_rs): + self.load_ckpt(_ck) + _state = torch.load(_rs, map_location=self.device, weights_only=False) + optimizer.load_state_dict(_state["optimizer"]) + if sched is not None and _state.get("sched") is not None: + sched.load_state_dict(_state["sched"]) + start_epoch = int(_state["epoch"]) + 1 + global_step = int(_state.get("global_step", 0)) + optimizer_step = int(_state.get("optimizer_step", 0)) + _mh = os.path.join(self.exp_path, "metrics_history.json") + if os.path.isfile(_mh): + with open(_mh) as f: + metrics_history = json.load(f)[:start_epoch] + if _is_main_process(): + logger.info(f"Resuming from epoch {start_epoch}/{epochs}") + elif os.path.isfile(_ck): + # Partial run with a weights checkpoint but no optimizer state + # (e.g. produced before atomic _resume.pt existed): load the + # trained weights and infer the epoch from metrics_history so we + # finish the remaining epochs instead of cold-restarting from 0. + self.load_ckpt(_ck) + _mh = os.path.join(self.exp_path, "metrics_history.json") + if os.path.isfile(_mh): + with open(_mh) as f: + metrics_history = json.load(f) + start_epoch = len(metrics_history) + if _is_main_process(): + logger.info(f"Warm-resuming from last.ckpt at epoch " + f"{start_epoch}/{epochs} (no optimizer state)") + + # Opt-in W&B tracking (full runs only: exp_path is None under Optuna's + # enable_logging=False, so per-trial training never spawns runs). + wrun = None + if self.exp_path is not None and _is_main_process(): + # Legible naming/config: exp dir basename is "{arch}_{method}_{task}_seed{N}". + _exp_name = getattr(self, "exp_name", None) or os.path.basename(self.exp_path.rstrip("/")) + _arch = _method = _task = _seed = None + if _exp_name and "_seed" in _exp_name: + _body, _, _seed = _exp_name.rpartition("_seed") + _parts = _body.split("_", 2) + if len(_parts) == 3: + _arch, _method, _task = _parts + _mdl = self.model.module if hasattr(self.model, "module") else self.model # unwrap DDP + wrun = _wandb.init_run( + config={"exp_name": _exp_name, "arch": _arch, "method": _method, + "task": _task, "seed": _seed, "epochs": epochs, + "lr": optimizer_params.get("lr"), "weight_decay": weight_decay, + "batch_size": getattr(train_dataloader, "batch_size", None), + "model": type(_mdl).__name__, "kind": "pretrain"}, + name=_exp_name, group=_task, job_type="pretrain", + tags=["kind:pretrain", "stage:pretrain", + f"bb:{_arch}" if _arch else None, + f"mod:{_task}" if _task else None, + f"method:{_method}" if _method else None]) + + epoch_iterator = tqdm( + range(start_epoch, epochs), + initial=start_epoch, total=epochs, + desc="Pretrain epochs", + unit="epoch", + disable=not _is_main_process(), + ) + for epoch in epoch_iterator: + epoch_iterator.set_postfix_str(f"{epoch + 1}/{epochs}", refresh=False) + if isinstance(train_dataloader.sampler, DistributedSampler): + train_dataloader.sampler.set_epoch(epoch) + + self.model.train() + if torch.cuda.is_available() and str(self.device).startswith("cuda"): + torch.cuda.reset_peak_memory_stats(self.device) + epoch_start = time.perf_counter() + epoch_losses = [] + epoch_loss_dicts: List[Dict[str, float]] = [] + + for _ in trange( + steps_per_epoch, + desc=f"Epoch {epoch + 1}/{epochs}", + smoothing=0.05, + leave=False, + disable=not _is_main_process(), + ): + try: + data = next(data_iterator) + except StopIteration: + data_iterator = iter(train_dataloader) + data = next(data_iterator) + + data = self._to_device(data) + + with torch.amp.autocast(device_type="cuda", enabled=use_amp): + output = self.model(**data) + loss = output["loss"] / grad_accumulation_steps + + if use_amp: + scaler.scale(loss).backward() + else: + loss.backward() + + if (global_step + 1) % grad_accumulation_steps == 0: + if max_grad_norm is not None: + if use_amp: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_( + self.model.parameters(), max_grad_norm + ) + + if use_amp: + # Detect whether GradScaler actually applied the step + # (it skips on inf/NaN grads, shrinking the scale). + scale_before = scaler.get_scale() + scaler.step(optimizer) + scaler.update() + stepped = scaler.get_scale() >= scale_before + else: + optimizer.step() + stepped = True + optimizer.zero_grad() + + # Only advance the schedule / EMA when a real update happened. + if stepped: + optimizer_step += 1 + # Warmup + cosine decay are handled by the LambdaLR + # scheduler, which steps once per optimizer step. + if sched is not None: + sched.step() + + if ( + self.ema_update_fn is not None + and optimizer_step % self.ema_update_every == 0 + ): + # Advance the EMA momentum schedule (cosine ->end) + # on the optimizer-step clock, then EMA-copy. + if hasattr(raw_model, "set_ema_decay"): + raw_model.set_ema_decay(optimizer_step, total_optim_steps) + # EMA function lives on the unwrapped model. + self.ema_update_fn() + + epoch_losses.append(loss.item() * grad_accumulation_steps) + if "loss_dict" in output: + epoch_loss_dicts.append(output["loss_dict"]) + + global_step += 1 + + epoch_time = time.perf_counter() - epoch_start + vram = _vram_stats(self.device) + + avg_loss_dict: Dict[str, float] = {} + if epoch_loss_dicts: + keys = set(k for d in epoch_loss_dicts for k in d) + for k in keys: + vals = [d[k] for d in epoch_loss_dicts if k in d] + if vals: + avg_loss_dict[k] = sum(vals) / len(vals) + + epoch_record: Dict[str, object] = { + "epoch": epoch, + "global_step": global_step, + "train_loss": sum(epoch_losses) / len(epoch_losses), + "epoch_time_s": round(epoch_time, 3), + "learning_rate": optimizer.param_groups[0]["lr"], + **avg_loss_dict, + **{f"train_{k}": v for k, v in vram.items()}, + } + + # Optional held-out SSL loss (used e.g. as the Optuna objective). + if val_dataloader is not None: + epoch_record.update(self._validate(val_dataloader, use_amp)) + + if _is_main_process(): + logger.info(f"--- Pretrain epoch-{epoch}, step-{global_step} ---") + logger.info(f"loss: {epoch_record['train_loss']:.4f}") + if "val_loss" in epoch_record: + logger.info(f"val_loss: {epoch_record['val_loss']:.4f}") + logger.info(f"epoch_time: {epoch_time:.2f}s") + if vram: + logger.info( + f"vram_peak: {vram['vram_peak_mb']:.1f} MB " + f"vram_current: {vram['vram_allocated_mb']:.1f} MB" + ) + + metrics_history.append(epoch_record) + _wandb.log(wrun, epoch_record, step=epoch) + + # Per-epoch hook (e.g. Optuna pruning); may raise to abort early. + if epoch_callback is not None: + epoch_callback(epoch, epoch_record) + + if _is_main_process() and self.exp_path is not None: + self.save_ckpt(os.path.join(self.exp_path, "last.ckpt")) + # resume state (optimizer/scheduler/epoch) — written atomically so + # a preemption mid-write can't corrupt it. + _rs = os.path.join(self.exp_path, "_resume.pt") + torch.save({"epoch": epoch, "global_step": global_step, + "optimizer_step": optimizer_step, + "optimizer": optimizer.state_dict(), + "sched": sched.state_dict() if sched is not None else None}, + _rs + ".tmp") + os.replace(_rs + ".tmp", _rs) + if (epoch + 1) % save_every_n_epochs == 0: + self.save_ckpt( + os.path.join(self.exp_path, f"epoch_{epoch + 1}.ckpt") + ) + history_path = os.path.join(self.exp_path, "metrics_history.json") + with open(history_path, "w") as f: + json.dump(metrics_history, f, indent=2) + + if self._ddp: + dist.barrier() + + if _is_main_process(): + total_time = time.perf_counter() - train_start + logger.info(f"--- Pretraining complete: {total_time:.2f}s total ---") + _wandb.finish(wrun) + + return metrics_history + + @torch.no_grad() + def _validate(self, val_dataloader: DataLoader, use_amp: bool = False) -> Dict[str, float]: + """Mean SSL loss (and per-component breakdown) over a held-out loader.""" + self.model.eval() + losses: List[float] = [] + loss_dicts: List[Dict[str, float]] = [] + for data in val_dataloader: + data = self._to_device(data) + with torch.amp.autocast(device_type="cuda", enabled=use_amp): + output = self.model(**data) + losses.append(output["loss"].item()) + if "loss_dict" in output: + loss_dicts.append(output["loss_dict"]) + self.model.train() + if not losses: + return {} + record: Dict[str, float] = {"val_loss": sum(losses) / len(losses)} + if loss_dicts: + keys = set(k for d in loss_dicts for k in d) + for k in keys: + vals = [d[k] for d in loss_dicts if k in d] + if vals: + record[f"val_{k}"] = sum(vals) / len(vals) + return record + + def _to_device(self, data): + """Recursively move tensors in a nested dict to the trainer device.""" + + def _move(obj): + if isinstance(obj, torch.Tensor): + return obj.to(self.device) + if isinstance(obj, dict): + return {k: _move(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return type(obj)(_move(x) for x in obj) + return obj + + return _move(data) + + def save_ckpt(self, ckpt_path: str) -> None: + """Save model state dict and training metadata.""" + Path(ckpt_path).parent.mkdir(parents=True, exist_ok=True) + raw_model = self.model.module if isinstance(self.model, DDP) else self.model + torch.save(raw_model.state_dict(), ckpt_path) + + def load_ckpt(self, ckpt_path: str) -> None: + """Load model state dict.""" + state_dict = torch.load( + ckpt_path, map_location=self.device, weights_only=True + ) + raw_model = self.model.module if isinstance(self.model, DDP) else self.model + raw_model.load_state_dict(state_dict) diff --git a/pyhealth/models/pretrain/utils.py b/pyhealth/models/pretrain/utils.py new file mode 100644 index 000000000..37eb47813 --- /dev/null +++ b/pyhealth/models/pretrain/utils.py @@ -0,0 +1,47 @@ +"""Utilities for wiring SSL pretraining models to PyHealth batches.""" + +from __future__ import annotations + +import torch + +from ...processors.base_processor import TemporalFeatureProcessor + + +def build_unified_inputs_from_batch( + processors: dict[str, TemporalFeatureProcessor], + feature_keys: list[str], + batch: dict[str, torch.Tensor | tuple[torch.Tensor, ...]], + device: torch.device | str | None = None, +) -> dict[str, dict[str, torch.Tensor]]: + """Convert a PyHealth collated batch into unified-embedding inputs. + + Each feature in ``batch`` is either a tensor or a tuple of tensors ordered + according to the processor's ``schema()``. This mirrors + :meth:`pyhealth.models.transformer.Transformer._build_unified_inputs`. + + Args: + processors: ``dataset.input_processors``. + feature_keys: Ordered list of input feature names. + batch: Collated batch dict from a PyHealth DataLoader. + device: If provided, move tensors to this device. + + Returns: + ``{field_name: {"value": Tensor, "time": Tensor, "mask": Tensor}}``. + """ + inputs: dict[str, dict[str, torch.Tensor]] = {} + for field_name in feature_keys: + feature = batch[field_name] + if isinstance(feature, torch.Tensor): + feature = (feature,) + schema = processors[field_name].schema() + field_dict: dict[str, torch.Tensor] = {} + if "value" in schema: + field_dict["value"] = feature[schema.index("value")] + if "time" in schema: + field_dict["time"] = feature[schema.index("time")] + if "mask" in schema: + field_dict["mask"] = feature[schema.index("mask")] + if device is not None: + field_dict = {k: v.to(device) for k, v in field_dict.items()} + inputs[field_name] = field_dict + return inputs diff --git a/scripts/kill_ddp.sh b/scripts/kill_ddp.sh new file mode 100644 index 000000000..2173f6a5f --- /dev/null +++ b/scripts/kill_ddp.sh @@ -0,0 +1,15 @@ +#!/bin/bash +# Hard-stop all full-pretrain DDP on this box, top-down (driver -> launcher -> +# torchrun -> workers), so torchrun can't restart workers mid-kill. +for pat in run_full_pretrain_local.sh run_full_pretrain.py "torch.distributed.run" torchrun "scripts/pretrain_ssl.py"; do + pkill -9 -f "$pat" +done +sleep 3 +# anything still holding a GPU, by PID +for p in $(nvidia-smi --query-compute-apps=pid --format=csv,noheader | sort -u); do + kill -9 "$p" 2>/dev/null +done +sleep 20 +echo "=== after kill ===" +nvidia-smi --query-gpu=index,memory.used --format=csv,noheader +echo "remaining workers: $(pgrep -f 'scripts/pretrain_ssl.py' | wc -l)" diff --git a/scripts/pretrain_ssl.py b/scripts/pretrain_ssl.py new file mode 100644 index 000000000..26a8f6431 --- /dev/null +++ b/scripts/pretrain_ssl.py @@ -0,0 +1,531 @@ +"""Self-supervised pretraining script for PyHealth multimodal sequences. + +Supports MAE, SimMIM, and I-JEPA over the unified embedding model. Run after +this finishes, use ``scripts/train_unified.py --pretrained-ckpt ...`` to +fine-tune downstream. + +Example: + python scripts/pretrain_ssl.py \ + --ehr-root /data/mimic-iv/2.2 \ + --note-root /data/mimic-iv/note \ + --task notes_labs \ + --method mae \ + --epochs 50 --batch-size 32 \ + --output-dir ./output/pretrain_mae +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path +from typing import Any, Dict + +import torch +import yaml + +# Make project root importable. +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from pyhealth.datasets import MIMIC4Dataset, get_dataloader +from pyhealth.models import UnifiedMultimodalEmbeddingModel +from pyhealth.models.pretrain import ( + MultimodalIJEPA, + MultimodalMaskedAutoencoder, + MultimodalSimMIM, + MultimodalVJEPA, +) +from pyhealth.models.pretrain.backbones import ARCH_CHOICES, build_backbone +from pyhealth.tasks.multimodal_mimic4 import ( + ClinicalNotesICDLabsMIMIC4, + ICDLabsMIMIC4, + LabsOnlyMIMIC4, + NotesLabsMIMIC4, +) +from pyhealth.models.pretrain.trainer import PretrainTrainer +from pyhealth.utils import set_seed + + +def _load_yaml(path: Path) -> Dict[str, Any]: + with open(path) as f: + return yaml.safe_load(f) or {} + + +def _resolve_config(config_path: Path) -> Dict[str, Any]: + """Load config and recursively merge _inherit chain (child wins).""" + cfg = _load_yaml(config_path) + inherit = cfg.pop("_inherit", None) + if inherit: + parent = _resolve_config(config_path.parent / inherit) + cfg = {**parent, **cfg} + return cfg + + +def _dict_to_namespace(cfg: Dict[str, Any], defaults: argparse.Namespace) -> argparse.Namespace: + """Convert a resolved config dict into an argparse Namespace. + + All CLI defaults are copied first so that optional config keys (e.g. + ``cache_dir``) always have an attribute. Config values then override + defaults. + """ + ns = argparse.Namespace(**vars(defaults)) + for k, v in cfg.items(): + setattr(ns, k, v) + return ns + + +def _merge_cli_overrides(ns: argparse.Namespace, cli: argparse.Namespace) -> argparse.Namespace: + """Explicit CLI-set values override config values.""" + for key in vars(cli): + val = getattr(cli, key) + if val is not None and key != "config": + setattr(ns, key, val) + return ns + + +# Fallback defaults applied after YAML config + CLI merging. argparse defaults +# are intentionally None so that config values are never clobbered by CLI defaults. +_DEFAULTS: Dict[str, Any] = { + "task": "labs_only", + "method": "mae", + # Encoder backbone architecture: transformer | jamba | mamba. + "arch": "transformer", + # Standardized compute: 128-dim, 2 layers, 4 heads (matches base.yaml and the + # downstream e2e backbone). These are only fallbacks; base.yaml provides them. + "embedding_dim": 128, + "heads": 4, + "num_layers": 2, + "dropout": 0.1, + # Mamba / Jamba backbone knobs. + "state_size": 16, + "conv_kernel": 4, + "jamba_transformer_layers": 1, + "jamba_mamba_layers": 1, + "rope_max_seq_len": 8192, + "rope_base": 10000.0, + "rope_scaling": 1.0, + "decoder_layers": 2, + "decoder_heads": 2, + "predictor_layers": 2, + "predictor_heads": 2, + "ema_decay": 0.996, + "ema_end": 1.0, + "num_target_blocks": 4, + "target_block_len": 4, + # V-JEPA: multi-scale spans + cross-modal windows. + "target_block_scales": [2, 4, 8], + "require_multimodal_blocks": False, + "normalize_targets": True, + # store_true flags: keep False fallbacks here and None argparse defaults so a + # config value of true is not clobbered by the absent-flag default. + "use_rope": False, + "norm_pix_loss": False, + "use_amp": False, + "icd_codes": False, + "include_vitals": False, + "freeze_encoder": False, + "mask_ratio": 0.5, + "mask_strategy": "random", + "epochs": 10, + "batch_size": 32, + "lr": 1e-4, + "weight_decay": 0.05, + "max_grad_norm": 1.0, + "scheduler": "cosine", + "warmup_steps": 1000, + "save_every_n_epochs": 5, + "num_workers": 4, + "seed": 42, + "grad_accumulation_steps": 1, + "local_rank": 0, + "observation_window_hours": 24, + "note_source": "discharge", + "note_extraction": "regex", + "text_finetune_mode": "full", + "dev": 0, + "output_dir": "./output/pretrain_ssl", +} + + +def _apply_defaults(ns: argparse.Namespace) -> argparse.Namespace: + for key, val in _DEFAULTS.items(): + if getattr(ns, key, None) is None: + setattr(ns, key, val) + return ns + + +def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset: + ehr_tables = ["diagnoses_icd", "procedures_icd", "labevents"] + note_tables = None + + if args.task == "clinical_notes_icd_labs": + if not args.note_root: + raise ValueError("--task clinical_notes_icd_labs requires --note-root.") + note_tables = ["discharge", "radiology"] + + 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")] + # notes_only reuses the same base dataset (tables) as notes_labs — only the + # task differs (include_labs=False), so the base-dataset cache is shared. + 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": + ehr_tables = ["labevents"] + note_tables = None + + 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, + cache_dir=args.cache_dir, + dev=args.dev if args.dev else False, + num_workers=args.num_workers, + ) + + +def _build_task(args: argparse.Namespace): + if args.task == "stagenet": + from pyhealth.tasks import MortalityPredictionStageNetMIMIC4 + return MortalityPredictionStageNetMIMIC4() + if args.task == "icd_labs": + 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 in ("notes_labs", "notes_only"): + task = NotesLabsMIMIC4( + 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"), + ) + 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}") + return task + if args.task == "labs_only": + return LabsOnlyMIMIC4(window_hours=args.observation_window_hours) + raise ValueError(f"Unknown task: {args.task}") + + +def _build_model(args: argparse.Namespace, sample_dataset: Any): + finetune_mode = "frozen" if args.freeze_encoder else args.text_finetune_mode + unified = UnifiedMultimodalEmbeddingModel( + processors=sample_dataset.input_processors, + embedding_dim=args.embedding_dim, + text_finetune_mode=finetune_mode, + ) + + backbone = build_backbone( + arch=getattr(args, "arch", "transformer"), + feature_size=args.embedding_dim, + num_layers=args.num_layers, + heads=args.heads, + dropout=args.dropout, + use_rope=getattr(args, "use_rope", False), + rope_max_seq_len=getattr(args, "rope_max_seq_len", 8192), + rope_base=getattr(args, "rope_base", 10000.0), + rope_scaling=getattr(args, "rope_scaling", 1.0), + state_size=getattr(args, "state_size", 16), + conv_kernel=getattr(args, "conv_kernel", 4), + num_transformer_layers=getattr(args, "jamba_transformer_layers", 1), + num_mamba_layers=getattr(args, "jamba_mamba_layers", 1), + ) + + per_modality_ratio = None + if args.lab_mask_ratio is not None or args.text_mask_ratio is not None: + per_modality_ratio = {} + # Map modality strings to indices using the unified model's lookup. + for field_name, modality in unified.modality_types.items(): + mod_idx = unified._modality_to_idx[modality] + if modality.value == "numeric" and args.lab_mask_ratio is not None: + per_modality_ratio[mod_idx] = args.lab_mask_ratio + if modality.value == "text" and args.text_mask_ratio is not None: + per_modality_ratio[mod_idx] = args.text_mask_ratio + + if args.method == "mae": + model = MultimodalMaskedAutoencoder( + embedding_model=unified, + backbone=backbone, + decoder_layers=args.decoder_layers, + decoder_heads=args.decoder_heads, + decoder_dim=args.decoder_dim, + mask_ratio=args.mask_ratio, + mask_strategy=args.mask_strategy, + per_modality_ratio=per_modality_ratio, + norm_pix_loss=args.norm_pix_loss, + ) + elif args.method == "simmim": + model = MultimodalSimMIM( + embedding_model=unified, + backbone=backbone, + mask_ratio=args.mask_ratio, + mask_strategy=args.mask_strategy, + per_modality_ratio=per_modality_ratio, + norm_targets=args.norm_pix_loss, + ) + elif args.method == "ijepa": + model = MultimodalIJEPA( + embedding_model=unified, + context_encoder=backbone, + predictor_layers=args.predictor_layers, + predictor_heads=args.predictor_heads, + predictor_dim=args.predictor_dim, + target_ema_decay=args.ema_decay, + target_ema_end=args.ema_end, + num_target_blocks=args.num_target_blocks, + target_block_len=args.target_block_len, + ) + elif args.method == "vjepa": + scales = getattr(args, "target_block_scales", None) or [2, 4, 8] + if isinstance(scales, str): + scales = [int(s) for s in scales.split(",") if s.strip()] + scales = tuple(int(s) for s in scales) + normalize_targets = getattr(args, "normalize_targets", None) + model = MultimodalVJEPA( + embedding_model=unified, + context_encoder=backbone, + predictor_layers=args.predictor_layers, + predictor_heads=args.predictor_heads, + predictor_dim=args.predictor_dim, + target_ema_decay=args.ema_decay, + target_ema_end=args.ema_end, + num_target_blocks=args.num_target_blocks, + target_block_scales=scales, + require_multimodal_blocks=getattr(args, "require_multimodal_blocks", False), + normalize_targets=True if normalize_targets is None else normalize_targets, + ) + else: + raise ValueError(f"Unknown pretraining method: {args.method}") + + # Attach dataset metadata so the model can convert raw batches internally. + model.feature_keys = list(sample_dataset.input_processors.keys()) + model.input_processors = sample_dataset.input_processors + return model + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description="SSL pretraining on unified multimodal sequences.") + p.add_argument("--config", type=str, default=None, help="YAML config path (e.g. configs/pretrain/mae_labs_only.yaml).") + p.add_argument("--ehr-root", type=str, default=None) + p.add_argument("--note-root", type=str, default=None) + p.add_argument("--cache-dir", type=str, default=None) + p.add_argument("--output-dir", type=str, default="./output/pretrain_ssl") + p.add_argument( + "--task", + type=str, + choices=["stagenet", "icd_labs", "clinical_notes_icd_labs", "notes_labs", "notes_only", "labs_only"], + default=None, + ) + p.add_argument( + "--method", + type=str, + choices=["mae", "simmim", "ijepa", "vjepa"], + default=None, + ) + + # Model + p.add_argument("--arch", type=str, default=None, choices=list(ARCH_CHOICES), + help="Encoder backbone architecture (transformer | jamba | mamba).") + p.add_argument("--embedding-dim", type=int, default=None) + p.add_argument("--heads", type=int, default=None) + p.add_argument("--num-layers", type=int, default=None) + p.add_argument("--dropout", type=float, default=None) + + # Mamba / Jamba backbone knobs + p.add_argument("--state-size", type=int, default=None, help="Mamba SSM state size (mamba/jamba).") + p.add_argument("--conv-kernel", type=int, default=None, help="Mamba causal conv kernel (mamba/jamba).") + p.add_argument("--jamba-transformer-layers", type=int, default=None, help="Attention layers in the Jamba stack.") + p.add_argument("--jamba-mamba-layers", type=int, default=None, help="Mamba layers in the Jamba stack.") + + # RoPE options + p.add_argument("--use-rope", action="store_true", default=None, help="Use RoPE in the Transformer backbone.") + p.add_argument("--rope-max-seq-len", type=int, default=None) + p.add_argument("--rope-base", type=float, default=None) + p.add_argument("--rope-scaling", type=float, default=None, help="NTK-aware scaling factor.") + + # MAE decoder + p.add_argument("--decoder-layers", type=int, default=None) + p.add_argument("--decoder-heads", type=int, default=None) + p.add_argument("--decoder-dim", type=int, default=None) + p.add_argument("--norm-pix-loss", action="store_true", default=None) + + # I-JEPA predictor + p.add_argument("--predictor-layers", type=int, default=None) + p.add_argument("--predictor-heads", type=int, default=None) + p.add_argument("--predictor-dim", type=int, default=None) + p.add_argument("--ema-decay", type=float, default=None) + p.add_argument("--ema-end", type=float, default=None) + p.add_argument("--num-target-blocks", type=int, default=None) + p.add_argument("--target-block-len", type=int, default=None) + + # V-JEPA specific + p.add_argument( + "--target-block-scales", + type=str, + default=None, + help="Comma-separated multi-scale block lengths for V-JEPA, e.g. '2,4,8'.", + ) + p.add_argument( + "--require-multimodal-blocks", + action="store_true", + default=None, + help="V-JEPA: prefer target windows that span more than one modality.", + ) + p.add_argument( + "--no-normalize-targets", + dest="normalize_targets", + action="store_false", + default=None, + help="V-JEPA: disable LayerNorm on EMA targets before the loss.", + ) + + # Masking + p.add_argument("--mask-ratio", type=float, default=None) + p.add_argument("--mask-strategy", type=str, default=None, choices=["random", "block"]) + p.add_argument("--lab-mask-ratio", type=float, default=None) + p.add_argument("--text-mask-ratio", type=float, default=None) + + # Training + p.add_argument("--epochs", type=int, default=None) + p.add_argument("--batch-size", type=int, default=None) + p.add_argument("--lr", type=float, default=None) + p.add_argument("--weight-decay", type=float, default=None) + p.add_argument("--max-grad-norm", type=float, default=None) + p.add_argument("--scheduler", type=str, default=None, choices=["none", "cosine"]) + p.add_argument("--warmup-steps", type=int, default=None) + p.add_argument("--save-every-n-epochs", type=int, default=None) + p.add_argument("--device", type=str, default=None) + p.add_argument("--num-workers", type=int, default=None) + p.add_argument("--seed", type=int, default=None) + p.add_argument("--grad-accumulation-steps", type=int, default=None) + p.add_argument("--use-amp", action="store_true", default=None, help="Use automatic mixed precision (CUDA only).") + + # torchrun / distributed + p.add_argument( + "--local-rank", + type=int, + default=0, + help="Local rank passed by torchrun (ignored, used for compatibility).", + ) + + # Task-specific + p.add_argument("--observation-window-hours", type=int, default=None) + p.add_argument("--icd-codes", action="store_true", default=None) + p.add_argument("--include-vitals", action="store_true", default=None) + p.add_argument("--note-source", type=str, default=None, choices=["discharge", "radiology"]) + p.add_argument( + "--note-extraction", + type=str, + default=None, + choices=[ + "regex", "regex_priority", "compact", "tfidf", + "section_hpi", "section_cc", "section_pmh", "section_meds", + "section_social", "section_family", "section_allergies", "section_ros", + "lab_retrieval", + ], + ) + p.add_argument("--tokenizer-model", type=str, default=None) + p.add_argument("--freeze-encoder", action="store_true", default=None) + p.add_argument( + "--text-finetune-mode", + type=str, + default=None, + help="frozen | full | topk:N | lora:r", + ) + + # Data + p.add_argument( + "--dev", + nargs="?", + type=int, + const=1000, + default=None, + help="Dev mode: limit dataset to N patients.", + ) + + return p.parse_args() + + +def main() -> None: + cli_args = parse_args() + + if cli_args.config: + config_path = Path(cli_args.config) + if not config_path.exists(): + raise SystemExit(f"Config not found: {config_path}") + cfg = _resolve_config(config_path) + args = _apply_defaults(_merge_cli_overrides(_dict_to_namespace(cfg, cli_args), cli_args)) + else: + args = _apply_defaults(cli_args) + + set_seed(args.seed) + + if not getattr(args, "ehr_root", None): + raise SystemExit("--ehr-root is required (either via CLI or config).") + + base_dataset = _build_base_dataset(args) + task = _build_task(args) + sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers) + + if len(sample_dataset) == 0: + raise RuntimeError("Task produced zero samples.") + + model = _build_model(args, sample_dataset) + + # EMA update hook for I-JEPA / V-JEPA (V-JEPA subclasses I-JEPA). + ema_fn = None + ema_every = 1 + if args.method in ("ijepa", "vjepa") and isinstance(model, MultimodalIJEPA): + ema_fn = model.update_target_encoder + + train_loader = get_dataloader(sample_dataset, batch_size=args.batch_size, shuffle=True) + + exp_name = f"{args.arch}_{args.method}_{args.task}_seed{args.seed}" + trainer = PretrainTrainer( + model=model, + device=args.device, + enable_logging=True, + output_path=args.output_dir, + exp_name=exp_name, + ema_update_fn=ema_fn, + ema_update_every=ema_every, + ) + + trainer.train( + train_dataloader=train_loader, + epochs=args.epochs, + optimizer_params={"lr": args.lr}, + weight_decay=args.weight_decay, + max_grad_norm=args.max_grad_norm, + scheduler=args.scheduler if args.scheduler != "none" else None, + warmup_steps=args.warmup_steps, + save_every_n_epochs=args.save_every_n_epochs, + grad_accumulation_steps=getattr(args, "grad_accumulation_steps", 1), + use_amp=getattr(args, "use_amp", False), + ) + + print(f"Pretraining complete. Logs: {trainer.exp_path}") + + +if __name__ == "__main__": + main() diff --git a/scripts/run_full_pretrain.py b/scripts/run_full_pretrain.py new file mode 100644 index 000000000..e17802b9d --- /dev/null +++ b/scripts/run_full_pretrain.py @@ -0,0 +1,102 @@ +"""Full-scale SSL pretraining at the tuned hyperparameters. + +Reads a ``best_params_pt___.json`` produced by +``optuna_pretrain.py`` and launches ``pretrain_ssl.py`` on the FULL dataset for +the real run (50 epochs), passing the tuned HPs as CLI overrides on top of the +128/2 base config. The resulting encoder checkpoint is what initializes the +downstream Table-2 runs. + +Example: + python scripts/run_full_pretrain.py \ + --best-params output/optuna_pretrain/notes_only/best_params_pt_mamba_vjepa_notes_only.json \ + --ehr-root ... --note-root ... --cache-dir ... --output-dir output/pretrain_full \ + --epochs 50 +""" +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] + +# best_params key -> pretrain_ssl.py flag. store_true flags handled separately. +_FLAG = { + "lr": "--lr", "weight_decay": "--weight-decay", "batch_size": "--batch-size", + "warmup_steps": "--warmup-steps", "max_grad_norm": "--max-grad-norm", + "mask_ratio": "--mask-ratio", "mask_strategy": "--mask-strategy", + "ema_decay": "--ema-decay", "num_target_blocks": "--num-target-blocks", + "state_size": "--state-size", "conv_kernel": "--conv-kernel", + "jamba_transformer_layers": "--jamba-transformer-layers", + "jamba_mamba_layers": "--jamba-mamba-layers", +} +_STORE_TRUE = {"norm_pix_loss": "--norm-pix-loss", "use_rope": "--use-rope"} + + +def build_cmd(cli) -> list: + meta = json.loads(Path(cli.best_params).read_text()) + arch, method, task = meta["arch"], meta["method"], meta["task"] + hp = meta["best_params"] + # task_only maps to the pretrain_ssl --task + flags (notes_only/labs_only/vitals) + task_flags = [] + if task == "notes_labs" and cli.include_vitals: + task_flags = ["--include-vitals"] + + args = ["--config", str(REPO_ROOT / "configs" / "pretrain" / "base.yaml"), + "--arch", arch, "--method", method, "--task", task, + "--ehr-root", cli.ehr_root, "--cache-dir", cli.cache_dir, + "--output-dir", cli.output_dir, "--epochs", str(cli.epochs), + "--num-workers", str(cli.num_workers), "--freeze-encoder", + *task_flags] + if cli.note_root: + args += ["--note-root", cli.note_root] + for k, flag in _FLAG.items(): + if k in hp and hp[k] is not None: + args += [flag, str(hp[k])] + for k, flag in _STORE_TRUE.items(): + if hp.get(k): + args.append(flag) + + # Appended last so they override the tuned values above (argparse keeps the + # final occurrence of a repeated flag). + args += [a for a in getattr(cli, "extra", []) or [] if a != "--"] + + script = str(REPO_ROOT / "scripts" / "pretrain_ssl.py") + if cli.nproc_per_node and cli.nproc_per_node > 1: + # multi-GPU DDP via torchrun (pretrain_ssl/PretrainTrainer are DDP-aware). + # Use the env's torchrun (next to sys.executable), not a bare PATH lookup + # that can resolve to the read-only system miniconda. + torchrun = str(Path(sys.executable).parent / "torchrun") + return [torchrun, "--standalone", f"--nproc_per_node={cli.nproc_per_node}", + script, *args] + return [sys.executable, script, *args] + + +def main(): + p = argparse.ArgumentParser(description="Full-scale SSL pretraining at tuned HPs.") + p.add_argument("--best-params", required=True) + p.add_argument("--ehr-root", required=True) + p.add_argument("--note-root", default=None) + p.add_argument("--cache-dir", required=True) + p.add_argument("--output-dir", required=True) + p.add_argument("--epochs", type=int, default=50) + p.add_argument("--num-workers", type=int, default=8) + p.add_argument("--nproc-per-node", type=int, default=1, help="GPUs for DDP (torchrun).") + p.add_argument("--include-vitals", action="store_true", default=False) + p.add_argument("--dry-run", action="store_true", default=False) + p.add_argument("--extra", nargs=argparse.REMAINDER, default=[], + help="Flags appended verbatim to pretrain_ssl.py, overriding the " + "tuned values (e.g. --batch-size 32 --grad-accumulation-steps 2 " + "to fit a smaller GPU at the same effective batch size).") + cli = p.parse_args() + cmd = build_cmd(cli) + print("[full-pretrain]", " ".join(cmd), flush=True) + if cli.dry_run: + return + sys.exit(subprocess.call(cmd)) + + +if __name__ == "__main__": + main() diff --git a/scripts/run_full_pretrain_local.sh b/scripts/run_full_pretrain_local.sh new file mode 100644 index 000000000..ada4c4195 --- /dev/null +++ b/scripts/run_full_pretrain_local.sh @@ -0,0 +1,47 @@ +#!/bin/bash +# Bare-metal full pretraining driver: ONE 4-GPU DDP job at a time (50% of an +# 8-GPU box). Runs a list of arch:method PAIRS sequentially on the 4 GPUs in +# GPUS. Resumable: skips runs whose last.ckpt already exists. Split a combo's 9 +# encoders across two boxes by giving each a different PAIRS subset. +# GPUS=0,1,2,3 COMBO=notes_labs PAIRS="transformer:mae jamba:vjepa" bash run_full_pretrain_local.sh +set -uo pipefail +source /home/rianatri/miniconda3/etc/profile.d/conda.sh && conda activate pyhealth2 +cd /home/rianatri/Multimodal-PyHealth-ssl +export PYTHONPATH=/home/rianatri/Multimodal-PyHealth-ssl TOKENIZERS_PARALLELISM=false PYTHONUNBUFFERED=1 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True # reduce fragmentation OOM (as in the slurm runner) +B=/shared/rsaas/rianatri/ssl +export HF_HOME=$B/huggingface TMPDIR=$B/tmp +export WANDB_PROJECT=pyhealth-multimodal WANDB_ENTITY=pyhealth-multimodal WANDB_DIR=$B/wandb + +COMBO=${COMBO:-notes_labs} +BP_DIR=${BP_DIR:-$B/optuna_pretrain} +GPUS=${GPUS:-0,1,2,3} +NGPU=$(echo "$GPUS" | tr "," "\n" | grep -c .) +# XTRA: extra pretrain_ssl.py flags appended last, overriding tuned values +# (e.g. XTRA="--batch-size 32 --grad-accumulation-steps 2" to fit a 48GB GPU at +# the same effective batch size). Must stay last: --extra is argparse REMAINDER. +case "$COMBO" in + notes_labs) TASK=notes_labs; EXTRA="" ;; + notes_labs_vitals) TASK=notes_labs; EXTRA="--include-vitals" ;; + *) TASK=$COMBO; EXTRA="" ;; +esac +OUT=$B/pretrain_full/$COMBO; mkdir -p "$OUT" "$B/logs" +EHR=/shared/rsaas/physionet.org/files/mimiciv/2.2 +NOTE=/shared/rsaas/physionet.org/files/mimic-note + +PAIRS_STR=${PAIRS:-"transformer:mae transformer:simmim transformer:vjepa jamba:mae jamba:simmim jamba:vjepa mamba:mae mamba:simmim mamba:vjepa"} +read -r -a PAIRS_ARR <<< "$PAIRS_STR" +for pair in "${PAIRS_ARR[@]}"; do + arch=${pair%:*}; method=${pair#*:} + bp=$BP_DIR/best_params_pt_${arch}_${method}_${TASK}.json + [ -f "$bp" ] || { echo "skip (no best_params): $arch $method"; continue; } + _mh="$OUT/${arch}_${method}_${TASK}_seed42/metrics_history.json" + [ -f "$_mh" ] && [ "$(python -c "import json;print(len(json.load(open('$_mh'))))" 2>/dev/null)" = "50" ] && { echo "skip (done, 50ep): $arch $method"; continue; } + echo "[$(date)] $arch/$method on GPUs $GPUS (ngpu=$NGPU)" + CUDA_VISIBLE_DEVICES=$GPUS python scripts/run_full_pretrain.py \ + --best-params "$bp" --ehr-root "$EHR" --note-root "$NOTE" --cache-dir "$B/cache_notes_labs" \ + --output-dir "$OUT" --epochs 50 --num-workers "${NW:-8}" --nproc-per-node "$NGPU" $EXTRA \ + ${XTRA:+--extra $XTRA} \ + > "$B/logs/fullpt_${COMBO}_${arch}_${method}.out" 2>&1 +done +echo "[$(date)] driver done (GPUS=$GPUS)." diff --git a/scripts/run_fullpt_condor.sh b/scripts/run_fullpt_condor.sh new file mode 100644 index 000000000..1943d7805 --- /dev/null +++ b/scripts/run_fullpt_condor.sh @@ -0,0 +1,8 @@ +#!/bin/bash +# Condor entrypoint for full pretraining on c02 at 4 GPUs (50% of the box). +# Condor sets CUDA_VISIBLE_DEVICES to the 4 GPUs it granted; hand those to the +# bare-metal driver so torchrun uses exactly them. COMBO/PAIRS come from the +# submit environment. +export GPUS="${CUDA_VISIBLE_DEVICES}" +cd /home/rianatri/Multimodal-PyHealth-ssl +exec bash scripts/run_full_pretrain_local.sh diff --git a/scripts/slurm/full_pretrain_cc.sh b/scripts/slurm/full_pretrain_cc.sh new file mode 100644 index 000000000..2c8a13af4 --- /dev/null +++ b/scripts/slurm/full_pretrain_cc.sh @@ -0,0 +1,49 @@ +#!/bin/bash +#SBATCH --job-name=ssl_full_pt +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=80G +#SBATCH --gres=gpu:1 +#SBATCH --requeue +#SBATCH --array=0-2 +#SBATCH --output=/scratch/rianatri/ssl/logs/full_pt_%x_%a_%j.out +# Full-scale (50-epoch) SSL pretraining at tuned HPs for one (combo, method), +# array over the 3 backbones. Partition/account/time set at submit time: +# V-JEPA -> IllinoisComputes-GPU A100 (-A jimeng-ic, 2-4 day) +# MAE/SimMIM -> eng-research-gpu A10 (-A jimeng-cs-eng) or scavenger +# FULL_COMBO in {labs_only, notes_only, notes_labs_vitals, notes_labs}; +# FULL_METHOD in {mae, simmim, vjepa}. +set -eo pipefail +source /scratch/rianatri/Multimodal-PyHealth-ssl/scripts/slurm/_env_cc.sh +[ -x "${PYBIN:-}/python" ] || { echo "FATAL: env not set up"; exit 1; } +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True # reduce fragmentation OOM + +ARCHS=( transformer jamba mamba ) +ARCH=${ARCHS[$SLURM_ARRAY_TASK_ID]} +COMBO=${FULL_COMBO:?set FULL_COMBO} +METHOD=${FULL_METHOD:?set FULL_METHOD} +case "${COMBO}" in + labs_only) TASK=labs_only; EXTRA="" ;; + notes_only) TASK=notes_only; EXTRA="" ;; + notes_labs_vitals) TASK=notes_labs; EXTRA="--include-vitals" ;; + *) TASK=notes_labs; EXTRA="" ;; +esac + +BP=${SSLBASE}/optuna_pretrain/${COMBO}/best_params_pt_${ARCH}_${METHOD}_${TASK}.json +[ -f "${BP}" ] || { echo "FATAL: best_params not found: ${BP}"; exit 1; } +OUT=${SSLBASE}/pretrain_full/${COMBO} +mkdir -p "${OUT}" +# resume-friendly: skip if a completed (50-epoch) encoder already exists +DONE="${OUT}/${ARCH}_${METHOD}_${TASK}_seed42/metrics_history.json" +if [ -f "${DONE}" ] && [ "$(${PYBIN}/python -c "import json;print(len(json.load(open('${DONE}'))))" 2>/dev/null)" = "50" ]; then + echo "skip (already 50 epochs): ${ARCH} ${METHOD} ${COMBO}"; exit 0 +fi +# GPUs allocated to this task -> DDP world size. +NGPU=$(echo "${CUDA_VISIBLE_DEVICES:-0}" | tr "," "\n" | grep -c .) +echo "[$(date)] full pretrain combo=${COMBO} method=${METHOD} arch=${ARCH} on $(hostname) ngpu=${NGPU}" +${PYBIN}/python scripts/run_full_pretrain.py \ + --best-params "${BP}" --ehr-root "${EHR_ROOT}" --note-root "${NOTE_ROOT}" \ + --cache-dir "${CACHE_DIR}" --output-dir "${OUT}" --epochs 50 --num-workers 8 \ + --nproc-per-node "${NGPU}" ${EXTRA} +echo "[$(date)] full pretrain combo=${COMBO} method=${METHOD} arch=${ARCH} done." diff --git a/tests/test_pretrain.py b/tests/test_pretrain.py new file mode 100644 index 000000000..685b9cb27 --- /dev/null +++ b/tests/test_pretrain.py @@ -0,0 +1,574 @@ +"""Unit tests for self-supervised pretraining models. + +Run with: + TOKENIZERS_PARALLELISM=false pytest tests/test_pretrain.py -v +""" + +import sys +import tempfile +from pathlib import Path + +import torch + + +def _make_code_dataset_and_batch(batch_size=2, seq_len=5): + """Build a minimal SampleDataset with one StageNetProcessor field.""" + from pyhealth.datasets import create_sample_dataset, get_dataloader + + codes_p0 = [f"c{i}" for i in range(seq_len)] + times_p0 = [float(i) for i in range(seq_len)] + codes_p1 = [f"c{i}" for i in range(2)] + times_p1 = [0.0, 1.0] + + samples = [ + { + "patient_id": "p0", + "visit_id": "v0", + "codes": (times_p0, codes_p0), + "label": 1, + }, + { + "patient_id": "p1", + "visit_id": "v1", + "codes": (times_p1, codes_p1), + "label": 0, + }, + ] + dataset = create_sample_dataset( + samples, + input_schema={"codes": "stagenet"}, + output_schema={"label": "binary"}, + dataset_name="test_pretrain", + ) + loader = get_dataloader(dataset, batch_size=batch_size, shuffle=False) + batch = next(iter(loader)) + return dataset, batch + + +def test_unified_mask_generator_random(): + from pyhealth.models.pretrain import UnifiedMaskGenerator + + mask = torch.tensor([[1, 1, 1, 1, 0], [1, 1, 0, 0, 0]]) + gen = UnifiedMaskGenerator(mask_ratio=0.5, strategy="random") + out = gen(mask) + assert out.shape == mask.shape + assert out.dtype == torch.bool + # Padding should never be masked. + assert not out[0, 4].item() + assert not out[1, 2:].any().item() + + +def test_unified_mask_generator_block(): + from pyhealth.models.pretrain import UnifiedMaskGenerator + + mask = torch.ones(2, 40) + gen = UnifiedMaskGenerator(mask_ratio=0.25, strategy="block", min_block_len=3, max_block_len=6) + out = gen(mask) + assert out.shape == mask.shape + # Roughly the requested ratio should be masked; exact count depends on + # block boundaries and trimming. + for b in range(out.shape[0]): + n_masked = out[b].sum().item() + assert 0 < n_masked <= int(mask.shape[1] * 0.5) + + +def test_mae_forward_loss(): + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.pretrain import MultimodalMaskedAutoencoder + from pyhealth.models.transformer import TransformerLayer + + dataset, batch = _make_code_dataset_and_batch() + unified = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, + embedding_dim=32, + ) + backbone = TransformerLayer(feature_size=32, heads=2, dropout=0.0, num_layers=1) + model = MultimodalMaskedAutoencoder( + embedding_model=unified, + backbone=backbone, + decoder_layers=1, + decoder_heads=2, + mask_ratio=0.5, + ) + model.feature_keys = list(dataset.input_processors.keys()) + model.input_processors = dataset.input_processors + + out = model(**batch) + assert "loss" in out + assert out["loss"].numel() == 1 + assert out["loss"].item() >= 0.0 + assert "pred" in out and "target" in out and "mask_token" in out + assert out["pred"].shape == out["target"].shape + out["loss"].backward() + + +def test_simmim_forward_loss(): + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.pretrain import MultimodalSimMIM + from pyhealth.models.transformer import TransformerLayer + + dataset, batch = _make_code_dataset_and_batch() + unified = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, + embedding_dim=32, + ) + backbone = TransformerLayer(feature_size=32, heads=2, dropout=0.0, num_layers=1) + model = MultimodalSimMIM( + embedding_model=unified, + backbone=backbone, + mask_ratio=0.5, + ) + model.feature_keys = list(dataset.input_processors.keys()) + model.input_processors = dataset.input_processors + + out = model(**batch) + assert "loss" in out + assert out["loss"].item() >= 0.0 + out["loss"].backward() + + +def test_ijepa_forward_loss(): + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.pretrain import MultimodalIJEPA + from pyhealth.models.transformer import TransformerLayer + + dataset, batch = _make_code_dataset_and_batch(seq_len=12) + unified = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, + embedding_dim=32, + ) + backbone = TransformerLayer(feature_size=32, heads=2, dropout=0.0, num_layers=1) + model = MultimodalIJEPA( + embedding_model=unified, + context_encoder=backbone, + predictor_layers=1, + predictor_heads=2, + predictor_dim=32, + num_target_blocks=2, + target_block_len=2, + ) + model.feature_keys = list(dataset.input_processors.keys()) + model.input_processors = dataset.input_processors + + out = model(**batch) + assert "loss" in out + assert out["loss"].item() >= 0.0 + # Target encoder should be an EMA copy and have no gradients. + for p in model.target_encoder.parameters(): + assert p.requires_grad is False + out["loss"].backward() + + +def test_rope_transformer_layer(): + from pyhealth.models.pretrain.rope import RoPETransformerLayer + + layer = RoPETransformerLayer(feature_size=32, heads=2, dropout=0.0, num_layers=1) + x = torch.randn(2, 10, 32) + mask = torch.ones(2, 10) + out, cls = layer(x, mask) + assert out.shape == (2, 10, 32) + assert cls.shape == (2, 32) + out.mean().backward() + + +def test_rope_extrapolation(): + from pyhealth.models.pretrain.rope import RotaryPositionEmbedding + + rope = RotaryPositionEmbedding(dim=16, max_seq_len=128, scaling_factor=2.0) + x = torch.randn(1, 200, 16) + out = rope(x, seq_len=200) + assert out.shape == x.shape + + +def test_load_pretrained_into_downstream_transformer(): + """Save a pretraining checkpoint and load it into a supervised Transformer.""" + import tempfile + from pyhealth.models import Transformer + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.pretrain import MultimodalSimMIM + from pyhealth.models.transformer import TransformerLayer + + dataset, batch = _make_code_dataset_and_batch() + unified = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, + embedding_dim=32, + ) + backbone = TransformerLayer(feature_size=32, heads=2, dropout=0.0, num_layers=1) + pretrain_model = MultimodalSimMIM( + embedding_model=unified, + backbone=backbone, + mask_ratio=0.5, + ) + pretrain_model.feature_keys = list(dataset.input_processors.keys()) + pretrain_model.input_processors = dataset.input_processors + + out = pretrain_model(**batch) + out["loss"].backward() + + with tempfile.NamedTemporaryFile(suffix=".ckpt", delete=False) as f: + ckpt_path = f.name + torch.save(pretrain_model.state_dict(), ckpt_path) + + # Build a supervised Transformer with the same unified embedding. + downstream = Transformer( + dataset=dataset, + embedding_dim=32, + heads=2, + num_layers=1, + unified_embedding=unified, + ) + + # Manually invoke the load helper from the e2e script. + import sys + sys.path.insert(0, str(Path(__file__).parent.parent / "examples" / "mortality_prediction")) + from unified_embedding_e2e_mimic4 import _load_pretrained_weights + + _load_pretrained_weights(downstream, ckpt_path) + + # Forward should still work. + out2 = downstream(**batch) + assert "loss" in out2 + out2["loss"].backward() + + import os + os.remove(ckpt_path) + + +def test_ijepa_blocks_not_dropped_and_distinct(): + """Regression: base I-JEPA must predict every target position (no dropped + blocks, no N x N broadcast) with distinct per-position predictions.""" + import torch + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.pretrain import MultimodalIJEPA + from pyhealth.models.transformer import TransformerLayer + + torch.manual_seed(0) + dataset, batch = _make_code_dataset_and_batch(seq_len=24) + unified = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, embedding_dim=32 + ) + backbone = TransformerLayer(feature_size=32, heads=2, dropout=0.0, num_layers=1) + model = MultimodalIJEPA( + embedding_model=unified, + context_encoder=backbone, + predictor_layers=1, + predictor_heads=2, + predictor_dim=32, + num_target_blocks=3, + target_block_len=2, + ) + model.feature_keys = list(dataset.input_processors.keys()) + model.input_processors = dataset.input_processors + model.eval() + + out = model(**batch) + pred = out["context_pred"] + tgt = out["target_embs"] + # Per-position (N, E), NOT an (N, N, E) broadcast. + assert pred.ndim == 2 and tgt.ndim == 2 + assert pred.shape == tgt.shape + # Number of predicted positions equals number of target positions (nothing + # silently dropped). + assert pred.shape[0] == int(out["target_mask"].sum().item()) + assert out["loss"].item() > 0.0 + import itertools + + dup = sum( + 1 + for i, j in itertools.combinations(range(pred.shape[0]), 2) + if torch.allclose(pred[i], pred[j], atol=1e-6) + ) + assert dup == 0 + + +def test_ijepa_single_block_nonzero_loss(): + """Regression: a single contiguous target block must still produce a real + (nonzero, gradient-bearing) loss, not the empty-fallback no-op.""" + import torch + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.pretrain import MultimodalIJEPA + from pyhealth.models.transformer import TransformerLayer + + torch.manual_seed(0) + dataset, batch = _make_code_dataset_and_batch(seq_len=16) + unified = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, embedding_dim=32 + ) + backbone = TransformerLayer(feature_size=32, heads=2, dropout=0.0, num_layers=1) + model = MultimodalIJEPA( + embedding_model=unified, + context_encoder=backbone, + predictor_layers=1, + predictor_heads=2, + predictor_dim=32, + num_target_blocks=1, + target_block_len=3, + ) + model.feature_keys = list(dataset.input_processors.keys()) + model.input_processors = dataset.input_processors + + out = model(**batch) + assert out["loss"].item() > 0.0 + out["loss"].backward() + # The predictor must have received gradient (loss is graph-connected). + assert any(p.grad is not None for p in model.predictor.parameters()) + + +def test_block_mask_preserves_contiguity(): + """Regression: block masking must yield a small number of contiguous runs, + not shattered single positions (the random-trim bug).""" + import torch + from pyhealth.models.pretrain import UnifiedMaskGenerator + + gen = UnifiedMaskGenerator(mask_ratio=0.5, strategy="block", min_block_len=3, max_block_len=12) + total_runs = 0 + trials = 30 + torch.manual_seed(0) + for _ in range(trials): + out = gen(torch.ones(1, 50))[0] + # Count contiguous runs of True. + runs = 0 + prev = False + for v in out.tolist(): + if v and not prev: + runs += 1 + prev = bool(v) + total_runs += runs + avg_runs = total_runs / trials + # With proper contiguous spans this is ~2-3; the buggy random-trim gave ~4.4+. + assert avg_runs < 4.0, f"avg contiguous runs {avg_runs} too high (spans shattered)" + + +def test_random_mask_floor(): + """Regression: every sample with valid positions masks at least one.""" + import torch + from pyhealth.models.pretrain import UnifiedMaskGenerator + + gen = UnifiedMaskGenerator(mask_ratio=0.3, strategy="random") + torch.manual_seed(0) + valid = torch.tensor([[1, 1, 1, 0, 0], [1, 1, 0, 0, 0], [1, 0, 0, 0, 0]]).float() + for _ in range(200): + out = gen(valid) + # No padding masked. + assert not (out & (valid == 0)).any() + # Every row with valid positions has >= 1 masked. + assert (out.any(dim=1) | (valid.sum(dim=1) == 0)).all() + + +def test_vjepa_forward_loss(): + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.pretrain import MultimodalVJEPA + from pyhealth.models.transformer import TransformerLayer + + dataset, batch = _make_code_dataset_and_batch(seq_len=20) + unified = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, + embedding_dim=32, + ) + backbone = TransformerLayer(feature_size=32, heads=2, dropout=0.0, num_layers=1) + model = MultimodalVJEPA( + embedding_model=unified, + context_encoder=backbone, + predictor_layers=1, + predictor_heads=2, + predictor_dim=32, + num_target_blocks=3, + target_block_scales=(2, 4), + ) + model.feature_keys = list(dataset.input_processors.keys()) + model.input_processors = dataset.input_processors + + out = model(**batch) + assert "loss" in out and out["loss"].item() >= 0.0 + assert "scale_ids" in out + # Target encoder must be a frozen EMA copy. + for p in model.target_encoder.parameters(): + assert p.requires_grad is False + out["loss"].backward() + # Predictor + scale embedding + context encoder receive gradients. + assert model.scale_embed.weight.grad is not None + assert any(p.grad is not None for p in model.context_encoder.parameters()) + # Target encoder receives NO gradient. + assert all(p.grad is None for p in model.target_encoder.parameters()) + + +def test_vjepa_predictor_distinguishes_blocks(): + """The V-JEPA fix: distinct target positions get distinct predictions. + + The base I-JEPA predictor used a single shared query, so multiple target + blocks in one sample produced identical predictions. V-JEPA's location + + scale aware queries must break that degeneracy. + """ + import torch + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.pretrain import MultimodalVJEPA + from pyhealth.models.transformer import TransformerLayer + + torch.manual_seed(0) + dataset, batch = _make_code_dataset_and_batch(seq_len=24) + unified = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, embedding_dim=32 + ) + backbone = TransformerLayer(feature_size=32, heads=2, dropout=0.0, num_layers=1) + model = MultimodalVJEPA( + embedding_model=unified, + context_encoder=backbone, + predictor_layers=1, + predictor_heads=2, + predictor_dim=32, + num_target_blocks=4, + target_block_scales=(2, 3), + ) + model.feature_keys = list(dataset.input_processors.keys()) + model.input_processors = dataset.input_processors + model.eval() + + out = model(**batch) + preds = out["context_pred"] # (N_target_positions, E) + assert preds.shape[0] >= 4 + # No two predicted positions should be exactly identical (degeneracy check). + import itertools + + dup = sum( + 1 + for i, j in itertools.combinations(range(preds.shape[0]), 2) + if torch.allclose(preds[i], preds[j], atol=1e-6) + ) + assert dup == 0, f"found {dup} identical predictions (predictor degenerate)" + + +def test_vjepa_multiscale_sampling(): + import torch + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.pretrain import MultimodalVJEPA + from pyhealth.models.transformer import TransformerLayer + + torch.manual_seed(1) + dataset, batch = _make_code_dataset_and_batch(seq_len=40) + unified = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, embedding_dim=16 + ) + backbone = TransformerLayer(feature_size=16, heads=2, dropout=0.0, num_layers=1) + model = MultimodalVJEPA( + embedding_model=unified, + context_encoder=backbone, + predictor_dim=16, + predictor_layers=1, + predictor_heads=2, + num_target_blocks=5, + target_block_scales=(2, 4, 8), + min_context_len=4, + ) + # Sample directly on a long all-valid sequence to exercise multi-scale draws. + event_mask = torch.ones(4, 40) + target_mask, context_mask, scale_ids = model._sample_multiscale_blocks(event_mask) + # Context always preserved. + assert (context_mask.sum(1) >= model.min_context_len).all() + # Targets and context are disjoint and cover only valid positions. + assert not (target_mask & context_mask).any() + # Over the batch, more than one scale index is used. + used_scales = scale_ids[target_mask].unique().numel() + assert used_scales >= 2 + + +def test_rope_inv_freq_spectrum(): + """Regression for the RoPE frequency bug (all freqs had collapsed to 1/base).""" + import torch + from pyhealth.models.pretrain.rope import RotaryPositionEmbedding + + dim = 32 + rope = RotaryPositionEmbedding(dim=dim, max_seq_len=128) + reference = 1.0 / (10000.0 ** (torch.arange(0, dim, 2).float() / dim)) + assert torch.allclose(rope.inv_freq, reference, atol=1e-6) + # Frequencies must be distinct (not collapsed to a single value). + assert rope.inv_freq.unique().numel() == dim // 2 + + +def test_ema_schedule_advances(): + """Regression: set_ema_decay must move the momentum toward target_ema_end.""" + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.pretrain import MultimodalIJEPA + from pyhealth.models.transformer import TransformerLayer + + dataset, _ = _make_code_dataset_and_batch() + unified = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, embedding_dim=16 + ) + backbone = TransformerLayer(feature_size=16, heads=2, dropout=0.0, num_layers=1) + model = MultimodalIJEPA( + embedding_model=unified, + context_encoder=backbone, + predictor_dim=16, + target_ema_decay=0.99, + target_ema_end=1.0, + ) + start = model.target_ema_decay + model.set_ema_decay(0, 100) + mid = model.target_ema_decay + model.set_ema_decay(100, 100) + end = model.target_ema_decay + assert abs(start - 0.99) < 1e-6 + assert mid >= start + assert abs(end - 1.0) < 1e-6 + + +def test_mae_simmim_token_target(): + """MAE/SimMIM default to the content-only ('token') target, not the + composed sequence (which leaks recoverable time/type into the loss).""" + import torch + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.pretrain import MultimodalMaskedAutoencoder, MultimodalSimMIM + from pyhealth.models.transformer import TransformerLayer + + dataset, batch = _make_code_dataset_and_batch(seq_len=8) + unified = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, embedding_dim=32 + ) + backbone = TransformerLayer(feature_size=32, heads=2, dropout=0.0, num_layers=1) + + for Model, kwargs in [ + (MultimodalMaskedAutoencoder, dict(decoder_layers=1, decoder_heads=2)), + (MultimodalSimMIM, dict()), + ]: + model = Model(embedding_model=unified, backbone=backbone, mask_ratio=0.5, **kwargs) + assert model.target == "token" # new default + model.feature_keys = list(dataset.input_processors.keys()) + model.input_processors = dataset.input_processors + + from pyhealth.models.pretrain.utils import build_unified_inputs_from_batch + inputs = build_unified_inputs_from_batch( + dataset.input_processors, list(dataset.input_processors.keys()), batch + ) + emb_out = unified(inputs) + assert "token_emb" in emb_out + # token_emb (content only) must differ from the composed sequence. + assert not torch.allclose(emb_out["token_emb"], emb_out["sequence"]) + + out = model(**batch) + # The returned target is the content-only token_emb, not the sequence. + assert torch.allclose(out["target"], emb_out["token_emb"], atol=1e-5) + assert not torch.allclose(out["target"], emb_out["sequence"]) + assert out["loss"].item() >= 0.0 + # Target MUST be detached (otherwise the model shrinks it -> collapse). + assert not out["target"].requires_grad + out["loss"].backward() + # But the embedding model must still train via the encoder INPUT path. + assert any( + p.grad is not None and p.grad.abs().sum() > 0 + for p in model.embedding_model.parameters() + ) + + +def test_per_modality_mae_decoder(): + from pyhealth.models.pretrain import PerModalityMAEDecoder + + decoder = PerModalityMAEDecoder( + embedding_dim=32, + output_specs={0: ("numeric", 10), 1: ("code", 50)}, + ) + emb = torch.randn(4, 8, 32) + type_ids = torch.tensor([[0, 0, 1, 1, 0, 1, 0, 1]]).expand(4, -1) + out = decoder(emb, type_ids) + assert 0 in out and 1 in out + assert out[0].shape == (out[0].shape[0], 10) + assert out[1].shape == (out[1].shape[0], 50) diff --git a/tests/test_pretrain_backbones.py b/tests/test_pretrain_backbones.py new file mode 100644 index 000000000..b805a2247 --- /dev/null +++ b/tests/test_pretrain_backbones.py @@ -0,0 +1,106 @@ +"""Tests for the SSL backbone factory (pyhealth/models/pretrain/backbones.py). + +Every backbone must satisfy the contract ``emb, cls = backbone(x, mask)`` with +``emb:(B,S,E)``, ``cls:(B,E)``, and padded positions must not contaminate the +outputs of valid positions. +""" +import pytest +import torch + +from pyhealth.models.pretrain.backbones import ARCH_CHOICES, build_backbone + +B, S, E = 3, 10, 16 + + +def _make(arch): + kw = dict(feature_size=E, num_layers=2, heads=4, dropout=0.1) + if arch == "jamba": + kw.update(num_transformer_layers=1, num_mamba_layers=1) + return build_backbone(arch, **kw) + + +@pytest.mark.parametrize("arch", ARCH_CHOICES) +def test_backbone_contract_shapes(arch): + m = _make(arch) + x = torch.randn(B, S, E) + mask = torch.ones(B, S) + mask[0, 6:] = 0 + emb, cls = m(x, mask) + assert emb.shape == (B, S, E) + assert cls.shape == (B, E) + + +@pytest.mark.parametrize("arch", ARCH_CHOICES) +def test_backbone_backward(arch): + m = _make(arch) + x = torch.randn(B, S, E, requires_grad=True) + emb, _ = m(x, torch.ones(B, S)) + emb.sum().backward() + assert x.grad is not None and torch.isfinite(x.grad).all() + + +@pytest.mark.parametrize("arch", ARCH_CHOICES) +def test_padding_does_not_leak_into_valid(arch): + """Perturbing padded positions must leave valid-position outputs unchanged.""" + m = _make(arch).eval() + x = torch.randn(B, S, E) + mask = torch.ones(B, S) + mask[0, 6:] = 0 + with torch.no_grad(): + emb, _ = m(x, mask) + x2 = x.clone() + x2[0, 6:] = 999.0 + emb2, _ = m(x2, mask) + assert torch.allclose(emb[0, :6], emb2[0, :6], atol=1e-5) + + +def test_unknown_arch_raises(): + with pytest.raises(ValueError): + build_backbone("gru", feature_size=E) + + +def test_transformer_rope_variant(): + m = build_backbone("transformer", feature_size=E, num_layers=1, heads=2, use_rope=True) + emb, cls = m(torch.randn(B, S, E), torch.ones(B, S)) + assert emb.shape == (B, S, E) and cls.shape == (B, E) + + +def _make_stagenet_dataset(n_codes=6): + from pyhealth.datasets import create_sample_dataset + + samples = [ + {"patient_id": "p0", "visit_id": "v0", + "codes": ([float(i) for i in range(n_codes)], [f"c{i}" for i in range(n_codes)]), "label": 1}, + {"patient_id": "p1", "visit_id": "v1", "codes": ([0.0, 1.0], ["c0", "c1"]), "label": 0}, + ] + return create_sample_dataset( + samples, input_schema={"codes": "stagenet"}, output_schema={"label": "binary"}, + dataset_name="test_pretrain_backbones", + ) + + +@pytest.mark.parametrize("arch", ARCH_CHOICES) +def test_mae_accepts_every_backbone(arch): + """Each backbone must plug into an SSL method and produce a finite, + differentiable loss with a per-modality loss_dict.""" + from pyhealth.datasets import get_dataloader + from pyhealth.models import UnifiedMultimodalEmbeddingModel + from pyhealth.models.pretrain import MultimodalMaskedAutoencoder + + dim = 32 + ds = _make_stagenet_dataset(n_codes=6) + batch = next(iter(get_dataloader(ds, batch_size=2, shuffle=False))) + unified = UnifiedMultimodalEmbeddingModel(processors=ds.input_processors, embedding_dim=dim) + kw = dict(feature_size=dim, num_layers=2, heads=4, dropout=0.0) + if arch == "jamba": + kw.update(num_transformer_layers=1, num_mamba_layers=1) + model = MultimodalMaskedAutoencoder( + embedding_model=unified, backbone=build_backbone(arch, **kw), + decoder_layers=2, decoder_heads=4, decoder_dim=dim, mask_ratio=0.5, + ) + model.feature_keys = list(ds.input_processors.keys()) + model.input_processors = ds.input_processors + out = model(**batch) + assert torch.isfinite(out["loss"]).item() + assert "total" in out["loss_dict"] + out["loss"].backward()