| """ |
| inference.py — Inference utilities for SleepMaMi pretrained / finetuned checkpoints. |
| |
| This module provides a thin, dependency-light wrapper around the two |
| classification backbones used in the paper and exposes three entry points: |
| |
| 1. extract_micro_features(...) |
| -> ClassificationMultiConv1d (epoch-level Micro-Encoder). |
| Returns the fused micro feature for each 30-s epoch. |
| |
| 2. extract_macro_micro_features(...) |
| -> ClassificationMultiConv1dMamba (full-night model). |
| Returns both the full-night Macro feature (sequence encoder output) |
| and the per-epoch Micro feature. |
| |
| 3. predict_sleep_stages(...) |
| -> ClassificationMultiConv1dMamba + classification head. |
| Returns the predicted sleep stage for every valid epoch of the night. |
| |
| Inputs are read from per-modality `.npy` files holding the full night as a |
| sequence of 30-s epochs, shape (N, sig_len) where sig_len = sig_len_target * |
| TARGET_RESAMPLE (3000 for the default 30 s @ 100 Hz). |
| |
| Example |
| ------- |
| # --checkpoint defaults to checkpoints/<task>.safetensors |
| python inference.py --task sleep_staging --npy_dir /path/to/night |
| |
| python inference.py --task micro --npy_dir /path/to/night --num_epochs 20 |
| |
| python inference.py --task macro_micro \ |
| --checkpoint checkpoints/macro_micro.safetensors \ |
| --npy_dir /path/to/night |
| """ |
|
|
| import argparse |
| import contextlib |
| import os |
|
|
| import numpy as np |
| import torch |
| from safetensors.torch import load_file as load_safetensors |
|
|
| import models |
| from configs.global_cfg import TARGET_RESAMPLE |
| from models.utils import flip_padded_sequences |
|
|
|
|
| |
| DEFAULT_MODALITIES = ["EEG", "EOG", "ECG", "RESP", "EMG", "SaO2"] |
|
|
|
|
| |
| |
| |
| def _load_state_dict(ckpt_path, map_location="cpu"): |
| """Return the raw parameter dict from a released `.safetensors` file |
| (weights only, keys without the `model.` prefix) or a Lightning / plain |
| `.ckpt` (e.g. one produced by `finetune.py`).""" |
| if ckpt_path.endswith(".safetensors"): |
| return load_safetensors(ckpt_path, device=str(map_location)) |
| |
| |
| ckpt = torch.load(ckpt_path, map_location=map_location, weights_only=False) |
| if isinstance(ckpt, dict) and "state_dict" in ckpt: |
| return ckpt["state_dict"] |
| return ckpt |
|
|
|
|
| def load_weights(model, ckpt_path, device="cpu", verbose=True): |
| """ |
| Load weights into `model`, handling the `model.` prefix that PyTorch |
| Lightning adds (the LightningModule stores the backbone as `self.model`). |
| |
| This mirrors `ConvMAE1d_Finetune.load_pretrained` in `finetune.py`: |
| position embeddings are skipped (they are buffers regenerated at build |
| time) and matching is attempted both with and without the `model.` prefix. |
| """ |
| src = _load_state_dict(ckpt_path, map_location=device) |
| tgt_keys = model.state_dict().keys() |
|
|
| new_state_dict = {} |
| for k in tgt_keys: |
| if "pos_embed" in k: |
| continue |
| if k in src: |
| new_state_dict[k] = src[k] |
| elif f"model.{k}" in src: |
| new_state_dict[k] = src[f"model.{k}"] |
|
|
| msg = model.load_state_dict(new_state_dict, strict=False) |
|
|
| if verbose: |
| missing = [k for k in msg.missing_keys if "pos_embed" not in k] |
| print(f"[load_weights] Loaded {len(new_state_dict)} tensors from {ckpt_path}") |
| if missing: |
| print(f"[load_weights] Missing (left at init): {missing}") |
| if msg.unexpected_keys: |
| print(f"[load_weights] Unexpected (ignored): {msg.unexpected_keys}") |
| return model |
|
|
|
|
| |
| |
| |
| def build_micro_model(modalities=DEFAULT_MODALITIES, sig_len_target=30, |
| patch_size=50, num_experts=4, num_experts_per=2, |
| num_classes=5, use_revin=True): |
| """ClassificationMultiConv1d — epoch-level Micro-Encoder (`classification_multi`).""" |
| return models.classification_multi( |
| modalities=modalities, |
| sig_len=sig_len_target * TARGET_RESAMPLE, |
| patch_size=patch_size, |
| num_experts=num_experts, |
| num_experts_per=num_experts_per, |
| num_classes=num_classes, |
| use_shared=True, |
| use_moe=True, |
| use_revin=use_revin, |
| ft_task="sleep_staging", |
| ) |
|
|
|
|
| def build_full_night_model(modalities=DEFAULT_MODALITIES, sig_len_target=30, |
| patch_size=50, num_experts=4, num_experts_per=2, |
| num_classes=5, use_revin=True, |
| use_micro=False, use_lstm=False, use_transformer=False, |
| bidi=True): |
| """ClassificationMultiConv1dMamba — full-night model (`classification_multi_full_night`).""" |
| return models.classification_multi_full_night( |
| modalities=modalities, |
| sig_len=sig_len_target * TARGET_RESAMPLE, |
| patch_size=patch_size, |
| num_experts=num_experts, |
| num_experts_per=num_experts_per, |
| num_classes=num_classes, |
| use_shared=True, |
| use_moe=True, |
| use_revin=use_revin, |
| ft_task="sleep_staging", |
| use_micro=use_micro, |
| use_lstm=use_lstm, |
| use_transformer=use_transformer, |
| bidi=False if use_transformer else bidi, |
| ) |
|
|
|
|
| |
| |
| |
| def _to_device(inputs, device, dtype): |
| return {m: v.to(device=device, dtype=dtype) for m, v in inputs.items()} |
|
|
|
|
| def _autocast(model, dtype=torch.bfloat16): |
| """Autocast context for the forward pass. |
| |
| FlashAttention (used in the encoder/fuser) only supports fp16/bf16, so a |
| pure float32 forward raises an error. Running under CUDA autocast keeps the |
| weights in their loaded dtype while computing attention in `dtype`, exactly |
| reproducing the `bf16-mixed` precision the checkpoints were trained with. |
| On CPU (no FlashAttention) autocast is a no-op. |
| """ |
| device = next(model.parameters()).device |
| if device.type == "cuda": |
| return torch.autocast(device_type="cuda", dtype=dtype) |
| return contextlib.nullcontext() |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| def _read_npy_arrays(npy_paths, modalities, sig_len): |
| """Load and shape per-modality arrays to (N, sig_len). Validates epoch counts.""" |
| arrays, n_epochs = {}, None |
| for m in modalities: |
| if m not in npy_paths: |
| raise ValueError(f"No .npy path provided for modality '{m}'") |
| arr = np.load(npy_paths[m]).astype(np.float32) |
| if arr.ndim == 1 or arr.shape[-1] != sig_len: |
| |
| |
| |
| |
| arr = arr.reshape(-1, sig_len) |
| if n_epochs is None: |
| n_epochs = arr.shape[0] |
| elif arr.shape[0] != n_epochs: |
| raise ValueError( |
| f"Modality '{m}' has {arr.shape[0]} epochs, expected {n_epochs} " |
| f"(all modalities must share the same number of epochs).") |
| arrays[m] = arr |
| return arrays, n_epochs |
|
|
|
|
| def _normalize_epoch(arr, modality): |
| """Per-modality normalization matching `data/data_loader.py`. |
| |
| - SaO2: scaled to [0, 1] via (clip(x, 0.7, 1.0) - 0.7) / 0.3. Percent-valued |
| inputs (max > 1.5) are divided by 100 first. The model downsamples SaO2 |
| internally, so it must be provided at the signal rate (sig_len samples). |
| - others: per-epoch z-normalization. With `use_revin=True` (default) the |
| model re-applies this internally, so the step is a harmless no-op there; |
| it matters when RevIN is disabled. |
| """ |
| if modality == "SaO2": |
| x = arr.copy() |
| if np.nanmax(x) > 1.5: |
| x = x / 100.0 |
| return (np.clip(x, 0.7, 1.0) - 0.7) / 0.3 |
| mu = arr.mean(axis=1, keepdims=True) |
| std = arr.std(axis=1, keepdims=True) |
| std = np.where(std < 1e-8, 1.0, std) |
| return (arr - mu) / std |
|
|
|
|
| def load_npy_epoch_inputs(npy_paths, modalities=DEFAULT_MODALITIES, sig_len_target=30, |
| normalize=True, num_epochs=None, device="cpu", |
| dtype=torch.float32): |
| """Epoch-level inputs for the Micro-Encoder, built from .npy files. |
| |
| Returns dict[str, Tensor] each of shape (N, sig_len, 1). |
| """ |
| sig_len = sig_len_target * TARGET_RESAMPLE |
| arrays, n = _read_npy_arrays(npy_paths, modalities, sig_len) |
| if num_epochs is not None: |
| n = min(num_epochs, n) |
| inputs = {} |
| for m in modalities: |
| a = arrays[m][:n] |
| if normalize: |
| a = _normalize_epoch(a, m) |
| inputs[m] = torch.from_numpy(np.ascontiguousarray(a)).unsqueeze(-1) |
| return _to_device(inputs, device, dtype) |
|
|
|
|
| def load_npy_full_night_inputs(npy_paths, modalities=DEFAULT_MODALITIES, sig_len_target=30, |
| normalize=True, device="cpu", dtype=torch.float32): |
| """Full-night inputs for the Mamba model, built from .npy files. |
| |
| Returns |
| ------- |
| inputs : dict[str, Tensor] each of shape (1, N*sig_len, 1) |
| mask : BoolTensor (1, N) all True — a single .npy night has no padding. |
| """ |
| sig_len = sig_len_target * TARGET_RESAMPLE |
| arrays, n = _read_npy_arrays(npy_paths, modalities, sig_len) |
| inputs = {} |
| for m in modalities: |
| a = arrays[m] |
| if normalize: |
| a = _normalize_epoch(a, m) |
| inputs[m] = torch.from_numpy(np.ascontiguousarray(a).reshape(1, n * sig_len, 1)) |
| mask = torch.ones(1, n, dtype=torch.bool, device=device) |
| return _to_device(inputs, device, dtype), mask |
|
|
|
|
| |
| |
| |
| @torch.inference_mode() |
| def extract_micro_features(model, inputs, amp_dtype=torch.bfloat16): |
| """ |
| (1) Micro feature from ClassificationMultiConv1d. |
| |
| Returns a tensor of shape (E, embed_dim[-1] * num_modalities) — one fused |
| micro embedding per 30-s epoch. |
| """ |
| model.eval() |
| with _autocast(model, amp_dtype): |
| feats = model.forward_features(inputs) |
| return feats |
|
|
|
|
| @torch.inference_mode() |
| def extract_macro_micro_features(model, inputs, mask, amp_dtype=torch.bfloat16): |
| """ |
| (2) Macro + Micro features from ClassificationMultiConv1dMamba. |
| |
| Returns |
| ------- |
| dict with |
| 'macro' : (1, num_epochs, macro_dim) full-night sequence feature |
| (for bidirectional models the forward and flipped-backward |
| streams are concatenated, exactly as in finetuning). |
| 'micro' : (1, num_epochs, micro_dim) per-epoch micro feature. |
| """ |
| model.eval() |
| with _autocast(model, amp_dtype): |
| macro, micro = model.forward_features(inputs, mask=mask) |
|
|
| if not model.use_micro and getattr(model, "bidi", False): |
| xf, xr = macro |
| xr = flip_padded_sequences(xr, mask) |
| macro = torch.cat([xf, xr], dim=-1) |
|
|
| return {"macro": macro, "micro": micro} |
|
|
|
|
| @torch.inference_mode() |
| def predict_sleep_stages(model, inputs, mask, return_logits=False, amp_dtype=torch.bfloat16): |
| """ |
| (3) Sleep-staging inference (Macro + Micro features + classification head). |
| |
| Returns |
| ------- |
| pred : LongTensor (num_valid_epochs,) predicted stage per valid epoch. |
| logits : (optional) FloatTensor (num_valid_epochs, num_classes). |
| """ |
| model.eval() |
| with _autocast(model, amp_dtype): |
| logits = model(inputs, mask=mask) |
| valid = mask[0].bool() |
| logits = logits[0][valid] |
| pred = logits.argmax(dim=-1) |
| if return_logits: |
| return pred, logits |
| return pred |
|
|
|
|
| |
| |
| |
| def _parse_args(): |
| p = argparse.ArgumentParser(description="SleepMaMi inference demo") |
| p.add_argument("--task", required=True, |
| choices=["micro", "macro_micro", "sleep_staging"], |
| help="Which inference to run.") |
| p.add_argument("--checkpoint", default=None, |
| help="Path to a .safetensors (released weights) or Lightning .ckpt file. " |
| "Default: checkpoints/<task>.safetensors") |
| |
| p.add_argument("--npy_dir", default=None, |
| help="Directory containing one '<MODALITY>.npy' per modality.") |
| p.add_argument("--npy", nargs="+", default=None, metavar="MOD=PATH", |
| help="Explicit per-modality .npy paths, e.g. EEG=eeg.npy EOG=eog.npy") |
| p.add_argument("--no_normalize", dest="normalize", action="store_false", |
| help="Skip built-in normalization (use if .npy is already preprocessed).") |
| p.set_defaults(normalize=True) |
| p.add_argument("--modalities", nargs="+", default=DEFAULT_MODALITIES) |
| p.add_argument("--sig_len_target", type=int, default=30) |
| p.add_argument("--patch_size", type=int, default=50) |
| p.add_argument("--num_classes", type=int, default=5) |
| p.add_argument("--num_experts", type=int, default=4) |
| p.add_argument("--num_experts_per", type=int, default=2) |
| p.add_argument("--no_revin", action="store_true", help="Disable RevIN normalization.") |
| p.add_argument("--num_epochs", type=int, default=None, |
| help="(micro only) number of epochs to embed; default = all.") |
| p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| p.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"]) |
| return p.parse_args() |
|
|
|
|
| def _resolve_npy_paths(args): |
| """Build a {modality: path} dict from --npy / --npy_dir.""" |
| if args.npy: |
| paths = {} |
| for item in args.npy: |
| if "=" not in item: |
| raise ValueError(f"--npy entries must be MOD=PATH, got '{item}'") |
| k, v = item.split("=", 1) |
| paths[k] = v |
| return paths |
| if args.npy_dir: |
| return {m: os.path.join(args.npy_dir, f"{m}.npy") for m in args.modalities} |
| raise SystemExit("Provide an input source: --npy_dir or --npy MOD=PATH ...") |
|
|
|
|
| def main(): |
| args = _parse_args() |
| if args.checkpoint is None: |
| args.checkpoint = os.path.join("checkpoints", f"{args.task}.safetensors") |
| device = args.device |
| dtype = {"float32": torch.float32, "bfloat16": torch.bfloat16, |
| "float16": torch.float16}[args.dtype] |
|
|
| npy_paths = _resolve_npy_paths(args) |
| common = dict(modalities=args.modalities, sig_len_target=args.sig_len_target, |
| patch_size=args.patch_size, num_experts=args.num_experts, |
| num_experts_per=args.num_experts_per, num_classes=args.num_classes, |
| use_revin=not args.no_revin) |
|
|
| if args.task == "micro": |
| model = build_micro_model(**common).to(device=device, dtype=dtype) |
| load_weights(model, args.checkpoint, device=device) |
| inputs = load_npy_epoch_inputs( |
| npy_paths, modalities=args.modalities, sig_len_target=args.sig_len_target, |
| normalize=args.normalize, num_epochs=args.num_epochs, |
| device=device, dtype=dtype) |
| feats = extract_micro_features(model, inputs) |
| print(f"Micro features: {tuple(feats.shape)} (epochs x feature_dim)") |
| print(feats.float().cpu().numpy()) |
|
|
| else: |
| model = build_full_night_model(**common).to(device=device, dtype=dtype) |
| load_weights(model, args.checkpoint, device=device) |
| inputs, mask = load_npy_full_night_inputs( |
| npy_paths, modalities=args.modalities, sig_len_target=args.sig_len_target, |
| normalize=args.normalize, device=device, dtype=dtype) |
|
|
| if args.task == "macro_micro": |
| out = extract_macro_micro_features(model, inputs, mask) |
| print(f"Macro feature: {tuple(out['macro'].shape)}") |
| print(f"Micro feature: {tuple(out['micro'].shape)}") |
| else: |
| pred = predict_sleep_stages(model, inputs, mask).cpu().numpy() |
| print(f"Predicted stages ({len(pred)} epochs): {pred}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
