diff --git a/README.md b/README.md index 50789be..707ddc5 100644 --- a/README.md +++ b/README.md @@ -68,9 +68,8 @@ python train_exposure_autoencoder.py \ --val_eid_file ukb_val_eid.csv ``` -The best checkpoint contains both `model_state_dict`, an `encoder_state_dict` -compatible with the default gated `TimesNetExposureEncoder`, and the channel -normalization statistics needed when the encoder is attached to DeepHealth. +The best checkpoint contains the encoder and normalization statistics needed +to generate fixed exposure embeddings. Multi-GPU pretraining follows the main trainer interface: add `--data_parallel --gpu_ids 0,1,2,3`. For efficient multi-GPU training, launch one process per GPU with DDP: @@ -84,13 +83,21 @@ torchrun --standalone --nproc_per_node=4 train_exposure_autoencoder.py \ In DDP mode, `--batch_size` is the global batch size and must be divisible by the number of processes. -The end-to-end next-step trainer supports the same DDP launch pattern: +Encode every cached exposure window once: + +```bash +python encode_exposure_cache.py \ + --exposure_cache_dir ukb_exposure_cache \ + --checkpoint runs/exposure_ae_RUN/best.pt +``` + +The next-step trainer reads `exposure_embeddings.npy` directly and does not +run TimesNet: ```bash torchrun --standalone --nproc_per_node=4 train_next_step.py \ --exposure_cache_dir ukb_exposure_cache \ - --batch_size 128 \ - --d_model 64 + --batch_size 128 ``` Training-channel statistics are cached at `/train_channel_stats.npz`; use diff --git a/dataset.py b/dataset.py index ac9d3ca..3b272f3 100644 --- a/dataset.py +++ b/dataset.py @@ -50,7 +50,11 @@ def _load_readonly_npy(path: Path) -> np.ndarray: class ExposureCache: """Eid-sequence-aligned exposure windows from prepare_exposure_cache.py.""" - def __init__(self, cache_dir: str | Path): + def __init__( + self, + cache_dir: str | Path, + embeddings_file: str | Path | None = None, + ): cache_dir = Path(cache_dir) self.cache_dir = cache_dir manifest_path = cache_dir / "exposure_manifest.json" @@ -101,6 +105,11 @@ class ExposureCache: self.eid_start = _load_readonly_npy(eid_start_path) self.daily = _load_readonly_npy(daily_path) self.monthly = _load_readonly_npy(monthly_path) + self.embeddings = ( + _load_readonly_npy(Path(embeddings_file)) + if embeddings_file is not None + else None + ) quality_path = cache_dir / "exposure_quality.npy" self.quality = _load_readonly_npy(quality_path) if quality_path.is_file() else None @@ -132,6 +141,16 @@ class ExposureCache: raise ValueError("exposure_eid_start.npy must have len(eid_index) + 1") if len(self.eid_start) and int(self.eid_start[-1]) != n_rows: raise ValueError("Last exposure eid offset must equal exposure row count") + if self.embeddings is not None: + if self.embeddings.ndim != 2: + raise ValueError( + "Exposure embeddings must have shape (N, D), got " + f"{self.embeddings.shape}" + ) + if max_window_index >= len(self.embeddings): + raise ValueError( + "Exposure row index points past exposure embeddings" + ) self._eid_to_pos = { int(eid): idx @@ -203,6 +222,20 @@ class ExposureCache: def monthly_windows(self, indices: np.ndarray) -> np.ndarray: return self._windows("monthly", indices) + def embedding_windows(self, indices: np.ndarray) -> np.ndarray: + if self.embeddings is None: + raise RuntimeError("Exposure embeddings are not enabled") + indices = np.asarray(indices, dtype=np.int64) + out = np.zeros( + (len(indices), self.embeddings.shape[1]), dtype=np.float32 + ) + valid_pos = np.nonzero(indices >= 0)[0] + if len(valid_pos): + out[valid_pos] = np.asarray( + self.embeddings[indices[valid_pos]], dtype=np.float32 + ) + return out + def _windows( self, kind: Literal["daily", "monthly"], @@ -288,18 +321,21 @@ class _ExpoBaseDataset(Dataset): no_event_interval_years: float = 5.0, include_no_event_in_uts_target: bool = False, exposure_cache_dir: str | Path | None = None, - mask_onset_exposure: bool = False, + exposure_embeddings_file: str | Path | None = None, ) -> None: self.data_prefix = data_prefix self.labels_file = labels_file self.no_event_interval_years = float(no_event_interval_years) self.include_no_event_in_uts_target = bool(include_no_event_in_uts_target) self.exposure_cache = ( - ExposureCache(exposure_cache_dir) + ExposureCache(exposure_cache_dir, exposure_embeddings_file) if exposure_cache_dir is not None else None ) - self.mask_onset_exposure = bool(mask_onset_exposure) + if exposure_embeddings_file is not None and exposure_cache_dir is None: + raise ValueError( + "exposure_cache_dir is required with exposure_embeddings_file" + ) self.label_code_to_id, self.label_id_to_code = load_label_vocab( labels_file, @@ -427,24 +463,14 @@ class _ExpoBaseDataset(Dataset): age_days=input_times_days, ) - def _load_exposure_windows(self, exposure_index: np.ndarray) -> tuple[torch.Tensor, torch.Tensor]: + def _load_exposure_embeddings( + self, exposure_index: np.ndarray + ) -> torch.Tensor: if self.exposure_cache is None: raise RuntimeError("Exposure cache is not enabled") - - daily = self.exposure_cache.daily_windows(exposure_index).astype( - np.float32, - copy=False, - ) - monthly = self.exposure_cache.monthly_windows(exposure_index).astype( - np.float32, - copy=False, - ) - - if self.mask_onset_exposure: - daily[:, 0, :] = np.nan - monthly[:, 0, :] = np.nan - - return torch.from_numpy(daily).float(), torch.from_numpy(monthly).float() + return torch.from_numpy( + self.exposure_cache.embedding_windows(exposure_index) + ).float() class NextStepHealthDataset(_ExpoBaseDataset): """ @@ -465,7 +491,7 @@ class NextStepHealthDataset(_ExpoBaseDataset): no_event_interval_years: float = 5.0, include_no_event_in_uts_target: bool = False, exposure_cache_dir: str | Path | None = None, - mask_onset_exposure: bool = False, + exposure_embeddings_file: str | Path | None = None, ) -> None: super().__init__( data_prefix=data_prefix, @@ -473,7 +499,7 @@ class NextStepHealthDataset(_ExpoBaseDataset): no_event_interval_years=no_event_interval_years, include_no_event_in_uts_target=include_no_event_in_uts_target, exposure_cache_dir=exposure_cache_dir, - mask_onset_exposure=mask_onset_exposure, + exposure_embeddings_file=exposure_embeddings_file, ) self.samples: List[Dict] = [] @@ -531,9 +557,9 @@ class NextStepHealthDataset(_ExpoBaseDataset): "target_multi_hot": torch.from_numpy(s["target_multi_hot"]).bool(), } if "exposure_index" in s: - daily, monthly = self._load_exposure_windows(s["exposure_index"]) - out["exposure_daily"] = daily - out["exposure_monthly"] = monthly + out["exposure_embedding"] = self._load_exposure_embeddings( + s["exposure_index"] + ) return out @@ -561,7 +587,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset): min_future_events: int = 1, validation_query_seed: int = 42, exposure_cache_dir: str | Path | None = None, - mask_onset_exposure: bool = False, + exposure_embeddings_file: str | Path | None = None, ) -> None: if split not in {"train", "valid", "test"}: raise ValueError(f"split must be train/valid/test, got {split!r}") @@ -572,7 +598,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset): no_event_interval_years=no_event_interval_years, include_no_event_in_uts_target=include_no_event_in_uts_target, exposure_cache_dir=exposure_cache_dir, - mask_onset_exposure=mask_onset_exposure, + exposure_embeddings_file=exposure_embeddings_file, ) self.split = split @@ -719,9 +745,9 @@ class AllFutureHealthDataset(_ExpoBaseDataset): input_times_days=times_days[hist], ) if exposure_index is not None: - daily, monthly = self._load_exposure_windows(exposure_index) - out["exposure_daily"] = daily - out["exposure_monthly"] = monthly + out["exposure_embedding"] = self._load_exposure_embeddings( + exposure_index + ) return out def __len__(self) -> int: @@ -745,19 +771,20 @@ def _collate_common_static(batch: List[Dict]) -> Dict: } -def _pad_exposure(batch: List[Dict], key: str, shape: tuple[int, int]) -> torch.Tensor: +def _pad_exposure_embedding(batch: List[Dict]) -> torch.Tensor: max_len = max(int(s["event_seq"].numel()) for s in batch) - out = torch.full( - (len(batch), max_len, shape[0], shape[1]), - float("nan"), - dtype=torch.float32, + embedding_dim = next( + int(s["exposure_embedding"].size(1)) + for s in batch + if "exposure_embedding" in s + ) + out = torch.zeros( + len(batch), max_len, embedding_dim, dtype=torch.float32 ) for idx, sample in enumerate(batch): - value = sample.get(key) - if value is None: - continue - seq_len = int(value.size(0)) - out[idx, :seq_len] = value + value = sample.get("exposure_embedding") + if value is not None: + out[idx, :value.size(0)] = value return out @@ -809,9 +836,8 @@ def next_step_collate_fn(batch: List[Dict]) -> Dict: "target_multi_hot": target_multi_hot, } out.update(_collate_common_static(batch)) - if any("exposure_daily" in s for s in batch): - out["exposure_daily"] = _pad_exposure(batch, "exposure_daily", DAILY_EXPOSURE_SHAPE) - out["exposure_monthly"] = _pad_exposure(batch, "exposure_monthly", MONTHLY_EXPOSURE_SHAPE) + if any("exposure_embedding" in s for s in batch): + out["exposure_embedding"] = _pad_exposure_embedding(batch) return out @@ -847,9 +873,8 @@ def all_future_collate_fn(batch: List[Dict]) -> Dict: "exposure": torch.stack([s["exposure"] for s in batch]), } out.update(_collate_common_static(batch)) - if any("exposure_daily" in s for s in batch): - out["exposure_daily"] = _pad_exposure(batch, "exposure_daily", DAILY_EXPOSURE_SHAPE) - out["exposure_monthly"] = _pad_exposure(batch, "exposure_monthly", MONTHLY_EXPOSURE_SHAPE) + if any("exposure_embedding" in s for s in batch): + out["exposure_embedding"] = _pad_exposure_embedding(batch) return out diff --git a/encode_exposure_cache.py b/encode_exposure_cache.py new file mode 100644 index 0000000..0c15e92 --- /dev/null +++ b/encode_exposure_cache.py @@ -0,0 +1,147 @@ +"""Encode cached exposure windows once for embedding-only model training.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import DataLoader, Dataset +from tqdm import tqdm + +from backbones import TimesNetExposureEncoder +from dataset import ExposureCache +from train_util import configure_torch_for_training, resolve_device + + +class ExposureEncodingDataset(Dataset): + def __init__(self, cache: ExposureCache): + self.cache = cache + + def __len__(self) -> int: + return len(self.cache.daily) + + def __getitem__(self, index: int) -> tuple[int, torch.Tensor, torch.Tensor]: + return ( + index, + torch.from_numpy( + np.array(self.cache.daily[index], dtype=np.float32, copy=True) + ), + torch.from_numpy( + np.array(self.cache.monthly[index], dtype=np.float32, copy=True) + ), + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Precompute exposure embeddings from an autoencoder checkpoint" + ) + parser.add_argument("--exposure_cache_dir", required=True) + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--output_file", default=None) + parser.add_argument("--batch_size", type=int, default=64) + parser.add_argument("--num_workers", type=int, default=4) + parser.add_argument("--device", default="cuda") + parser.add_argument( + "--output_dtype", choices=["float16", "float32"], default="float16" + ) + parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + if args.batch_size <= 0: + parser.error("--batch_size must be positive") + return args + + +def main() -> None: + args = parse_args() + device = resolve_device(args.device) + configure_torch_for_training(device) + cache_dir = Path(args.exposure_cache_dir) + output_path = ( + Path(args.output_file) + if args.output_file + else cache_dir / "exposure_embeddings.npy" + ) + if output_path.exists() and not args.overwrite: + raise FileExistsError(f"{output_path} exists; pass --overwrite") + + checkpoint_data = torch.load(args.checkpoint, map_location="cpu") + model_cfg = checkpoint_data["model_config"] + normalization = checkpoint_data["normalization"] + encoder = TimesNetExposureEncoder( + n_embd=int(model_cfg["n_embd"]), + d_model=model_cfg["d_model"], + n_layers=int(model_cfg["n_layers"]), + top_k=int(model_cfg["top_k"]), + n_backbone_blocks=int(model_cfg["n_backbone_blocks"]), + backbone_kernel_size=int(model_cfg["backbone_kernel_size"]), + backbone_expansion=float(model_cfg["backbone_expansion"]), + dropout=float(model_cfg["dropout"]), + use_gate=True, + ) + encoder.load_state_dict(checkpoint_data["encoder_state_dict"], strict=True) + encoder.to(device).eval() + + cache = ExposureCache(cache_dir) + loader_kwargs = { + "batch_size": args.batch_size, + "shuffle": False, + "num_workers": args.num_workers, + "pin_memory": device.type == "cuda", + "persistent_workers": args.num_workers > 0, + } + loader = DataLoader(ExposureEncodingDataset(cache), **loader_kwargs) + output_dtype = np.float16 if args.output_dtype == "float16" else np.float32 + output_path.parent.mkdir(parents=True, exist_ok=True) + output = np.lib.format.open_memmap( + output_path, + mode="w+", + dtype=output_dtype, + shape=(len(cache.daily), int(model_cfg["n_embd"])), + ) + stats = { + key: torch.as_tensor(normalization[key], device=device).view(1, 1, -1) + for key in ("daily_mean", "daily_std", "monthly_mean", "monthly_std") + } + + with torch.inference_mode(): + for indices, daily, monthly in tqdm(loader, desc="Encoding exposure"): + daily = daily.to(device, non_blocking=True) + monthly = monthly.to(device, non_blocking=True) + daily_mask = torch.isfinite(daily) + monthly_mask = torch.isfinite(monthly) + has_observation = ( + daily_mask.flatten(1).any(dim=1) + | monthly_mask.flatten(1).any(dim=1) + ) + daily = ( + torch.nan_to_num(daily) - stats["daily_mean"] + ) / stats["daily_std"] + monthly = ( + torch.nan_to_num(monthly) - stats["monthly_mean"] + ) / stats["monthly_std"] + daily = daily * daily_mask + monthly = monthly * monthly_mask + encoded = encoder(daily, monthly, daily_mask, monthly_mask) + encoded = encoded * has_observation.unsqueeze(1) + output[indices.numpy()] = encoded.float().cpu().numpy().astype( + output_dtype, copy=False + ) + output.flush() + metadata = { + "checkpoint": str(Path(args.checkpoint).resolve()), + "output_file": str(output_path.resolve()), + "rows": len(cache.daily), + "embedding_dim": int(model_cfg["n_embd"]), + "dtype": args.output_dtype, + } + output_path.with_suffix(".json").write_text( + json.dumps(metadata, indent=2), encoding="utf-8" + ) + print(f"Wrote {len(cache.daily):,} embeddings to {output_path}") + + +if __name__ == "__main__": + main() diff --git a/eval_data.py b/eval_data.py index af585c8..4111d33 100644 --- a/eval_data.py +++ b/eval_data.py @@ -8,8 +8,6 @@ import torch from torch.nn.utils.rnn import pad_sequence from dataset import ( - DAILY_EXPOSURE_SHAPE, - MONTHLY_EXPOSURE_SHAPE, AllFutureHealthDataset, HealthDataset, ) @@ -32,7 +30,7 @@ class AllFutureSequenceEvalDataset: min_history_events: int = 1, min_future_events: int = 1, exposure_cache_dir: str | Path | None = None, - mask_onset_exposure: bool = False, + exposure_embeddings_file: str | Path | None = None, ) -> None: base = AllFutureHealthDataset( data_prefix=data_prefix, @@ -41,7 +39,7 @@ class AllFutureSequenceEvalDataset: min_history_events=min_history_events, min_future_events=min_future_events, exposure_cache_dir=exposure_cache_dir, - mask_onset_exposure=mask_onset_exposure, + exposure_embeddings_file=exposure_embeddings_file, ) self.base = base @@ -91,9 +89,9 @@ class AllFutureSequenceEvalDataset: "sex": torch.tensor(s["sex"], dtype=torch.long), } if "exposure_index" in s: - daily, monthly = self.base._load_exposure_windows(s["exposure_index"]) - out["exposure_daily"] = daily - out["exposure_monthly"] = monthly + out["exposure_embedding"] = self.base._load_exposure_embeddings( + s["exposure_index"] + ) return out @@ -107,7 +105,7 @@ def load_sequence_eval_dataset( min_history_events: int, min_future_events: int, exposure_cache_dir: str | Path | None = None, - mask_onset_exposure: bool = False, + exposure_embeddings_file: str | Path | None = None, ): mode = str(model_target_mode).lower() if mode == "next_token": @@ -117,7 +115,7 @@ def load_sequence_eval_dataset( no_event_interval_years=no_event_interval_years, include_no_event_in_uts_target=include_no_event_in_uts_target, exposure_cache_dir=exposure_cache_dir, - mask_onset_exposure=mask_onset_exposure, + exposure_embeddings_file=exposure_embeddings_file, ) if mode == "all_future": return AllFutureSequenceEvalDataset( @@ -126,7 +124,7 @@ def load_sequence_eval_dataset( min_history_events=min_history_events, min_future_events=min_future_events, exposure_cache_dir=exposure_cache_dir, - mask_onset_exposure=mask_onset_exposure, + exposure_embeddings_file=exposure_embeddings_file, ) raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}") @@ -156,34 +154,17 @@ def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, "readout_mask": readout_mask, "sex": torch.stack([s["sex"] for s in batch]), } - if any("exposure_daily" in s for s in batch): - out["exposure_daily"] = _pad_eval_exposure( - batch, - "exposure_daily", - DAILY_EXPOSURE_SHAPE, + if any("exposure_embedding" in s for s in batch): + embedding_dim = next( + int(s["exposure_embedding"].size(1)) + for s in batch if "exposure_embedding" in s ) - out["exposure_monthly"] = _pad_eval_exposure( - batch, - "exposure_monthly", - MONTHLY_EXPOSURE_SHAPE, + encoded = torch.zeros( + len(batch), event_seq.size(1), embedding_dim, dtype=torch.float32 ) - return out - - -def _pad_eval_exposure( - batch: List[Dict[str, torch.Tensor]], - key: str, - shape: tuple[int, int], -) -> torch.Tensor: - max_len = max(int(s["event_seq"].numel()) for s in batch) - out = torch.full( - (len(batch), max_len, shape[0], shape[1]), - float("nan"), - dtype=torch.float32, - ) - for idx, sample in enumerate(batch): - value = sample.get(key) - if value is None: - continue - out[idx, : int(value.size(0))] = value + for idx, sample in enumerate(batch): + value = sample.get("exposure_embedding") + if value is not None: + encoded[idx, :value.size(0)] = value + out["exposure_embedding"] = encoded return out diff --git a/evaluate_auc.py b/evaluate_auc.py index bcb6f71..926cb39 100644 --- a/evaluate_auc.py +++ b/evaluate_auc.py @@ -323,14 +323,9 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data target_mode=model_target_mode, dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")), dropout=float(cfg_get(args, cfg, "dropout", 0.0)), - use_exposure_encoder=bool(cfg_get(args, cfg, "use_exposure_encoder", False)), - exposure_d_model=cfg_get(args, cfg, "d_model", 64), - exposure_n_layers=int(cfg_get(args, cfg, "exposure_n_layers", 2)), - exposure_top_k=int(cfg_get(args, cfg, "exposure_top_k", 2)), - exposure_n_backbone_blocks=int(cfg_get(args, cfg, "exposure_n_backbone_blocks", 1)), - exposure_backbone_kernel_size=int(cfg_get(args, cfg, "exposure_backbone_kernel_size", 5)), - exposure_backbone_expansion=float(cfg_get(args, cfg, "exposure_backbone_expansion", 2.0)), - exposure_use_gate=bool(cfg_get(args, cfg, "exposure_use_gate", True)), + use_exposure_embeddings=bool( + cfg_get(args, cfg, "use_exposure_embeddings", False) + ), ) @@ -531,14 +526,9 @@ def infer_readout_hidden( sex=batch_dev["sex"][active], padding_mask=padding_mask[active], t_query=time_seq[active, pos], - exposure_daily=( - batch_dev["exposure_daily"][active] - if "exposure_daily" in batch_dev - else None - ), - exposure_monthly=( - batch_dev["exposure_monthly"][active] - if "exposure_monthly" in batch_dev + exposure_embedding=( + batch_dev["exposure_embedding"][active] + if "exposure_embedding" in batch_dev else None ), target_mode="all_future", @@ -551,8 +541,7 @@ def infer_readout_hidden( time_seq=time_seq, sex=batch_dev["sex"], padding_mask=padding_mask, - exposure_daily=batch_dev.get("exposure_daily"), - exposure_monthly=batch_dev.get("exposure_monthly"), + exposure_embedding=batch_dev.get("exposure_embedding"), target_mode="next_token", ) ro = readout( @@ -1337,7 +1326,7 @@ def main() -> None: min_history_events=int(cfg.get("all_future_min_history_events", 1)), min_future_events=int(cfg.get("all_future_min_future_events", 1)), exposure_cache_dir=cfg.get("exposure_cache_dir", None), - mask_onset_exposure=bool(cfg.get("mask_onset_exposure", False)), + exposure_embeddings_file=cfg.get("exposure_embeddings_file", None), ) validate_dataset_metadata(dataset, cfg) diff --git a/models.py b/models.py index ed5b4d6..3b79f45 100644 --- a/models.py +++ b/models.py @@ -7,7 +7,6 @@ import torch.nn.functional as F from backbones import ( AgeSinusoidalEncoding, GPTBlock, - TimesNetExposureEncoder, ) from targets import PAD_IDX @@ -30,16 +29,7 @@ class DeepHealth(nn.Module): target_mode: str = "next_token", # "next_token" or "all_future" dist_mode: str = "exponential", # "exponential", "weibull" or "mixed" dropout: float = 0.0, - use_exposure_encoder: bool = False, - exposure_daily_input_dim: int = 4, - exposure_monthly_input_dim: int = 2, - exposure_d_model: int | None = 64, - exposure_n_layers: int = 2, - exposure_top_k: int = 2, - exposure_n_backbone_blocks: int = 1, - exposure_backbone_kernel_size: int = 5, - exposure_backbone_expansion: float = 2.0, - exposure_use_gate: bool = True, + use_exposure_embeddings: bool = False, ): super().__init__() if target_mode not in ["next_token", "all_future"]: @@ -53,26 +43,9 @@ class DeepHealth(nn.Module): 2, n_embd) # Assuming binary gender self.target_mode = target_mode self.dist_mode = dist_mode - self.use_exposure_encoder = use_exposure_encoder + self.use_exposure_embeddings = bool(use_exposure_embeddings) self.n_embd = n_embd self.vocab_size = vocab_size - self.exposure_encoder = ( - TimesNetExposureEncoder( - n_embd=n_embd, - daily_input_dim=exposure_daily_input_dim, - monthly_input_dim=exposure_monthly_input_dim, - d_model=exposure_d_model, - n_layers=exposure_n_layers, - top_k=exposure_top_k, - n_backbone_blocks=exposure_n_backbone_blocks, - backbone_kernel_size=exposure_backbone_kernel_size, - backbone_expansion=exposure_backbone_expansion, - dropout=dropout, - use_gate=exposure_use_gate, - ) - if use_exposure_encoder - else None - ) nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02) nn.init.zeros_(self.token_embedding.weight[0]) nn.init.normal_(self.gender_embedding.weight, mean=0.0, std=0.02) @@ -121,95 +94,6 @@ class DeepHealth(nn.Module): dtype=dtype, ).masked_fill(~valid, -1e4)[:, None, :, :] - def _encode_event_exposure( - self, - exposure_daily: torch.Tensor | None, - exposure_monthly: torch.Tensor | None, - exposure_daily_mask: torch.Tensor | None, - exposure_monthly_mask: torch.Tensor | None, - event_shape: tuple[int, int], - ) -> torch.Tensor | None: - if self.exposure_encoder is None: - if exposure_daily is not None or exposure_monthly is not None: - raise ValueError( - "Exposure tensors were provided but use_exposure_encoder=False" - ) - return None - if exposure_daily is None or exposure_monthly is None: - raise ValueError( - "exposure_daily and exposure_monthly are required when " - "use_exposure_encoder=True" - ) - - batch_size, event_len = event_shape - if exposure_daily.shape[:2] != event_shape: - raise ValueError( - "exposure_daily must have shape (B, L, T_daily, C_daily), got " - f"{tuple(exposure_daily.shape)} for event shape {event_shape}" - ) - if exposure_monthly.shape[:2] != event_shape: - raise ValueError( - "exposure_monthly must have shape (B, L, T_monthly, C_monthly), got " - f"{tuple(exposure_monthly.shape)} for event shape {event_shape}" - ) - if exposure_daily.dim() != 4 or exposure_monthly.dim() != 4: - raise ValueError( - "exposure_daily and exposure_monthly must both be 4D tensors" - ) - - def flatten_mask(mask: torch.Tensor | None, ref: torch.Tensor) -> torch.Tensor | None: - if mask is None: - return None - if mask.shape[:2] != event_shape: - raise ValueError( - "exposure mask must start with shape (B, L), got " - f"{tuple(mask.shape)}" - ) - if mask.dim() not in {3, 4}: - raise ValueError( - "exposure mask must have shape (B, L, T) or (B, L, T, C), " - f"got {tuple(mask.shape)}" - ) - if mask.shape[2] != ref.shape[2]: - raise ValueError( - "exposure mask time dimension does not match exposure tensor" - ) - if mask.dim() == 4 and mask.shape[3] != ref.shape[3]: - raise ValueError( - "exposure mask channel dimension does not match exposure tensor" - ) - return mask.reshape(batch_size * event_len, *mask.shape[2:]) - - daily = exposure_daily.reshape( - batch_size * event_len, - exposure_daily.size(2), - exposure_daily.size(3), - ) - monthly = exposure_monthly.reshape( - batch_size * event_len, - exposure_monthly.size(2), - exposure_monthly.size(3), - ) - daily_mask = flatten_mask(exposure_daily_mask, exposure_daily) - monthly_mask = flatten_mask(exposure_monthly_mask, exposure_monthly) - - exposure_device = exposure_daily.device - exposure_dtype = self.token_embedding.weight.dtype - daily = daily.to(device=exposure_device, dtype=exposure_dtype) - monthly = monthly.to(device=exposure_device, dtype=exposure_dtype) - if daily_mask is not None: - daily_mask = daily_mask.to(device=exposure_device) - if monthly_mask is not None: - monthly_mask = monthly_mask.to(device=exposure_device) - - exposure_emb = self.exposure_encoder( - daily=daily, - monthly=monthly, - daily_mask=daily_mask, - monthly_mask=monthly_mask, - ) - return exposure_emb.reshape(batch_size, event_len, self.n_embd) - def _forward_shared( self, event_seq: torch.LongTensor, @@ -218,10 +102,7 @@ class DeepHealth(nn.Module): mode: str, padding_mask: torch.Tensor | None = None, t_query: torch.FloatTensor | None = None, - exposure_daily: torch.Tensor | None = None, - exposure_monthly: torch.Tensor | None = None, - exposure_daily_mask: torch.Tensor | None = None, - exposure_monthly_mask: torch.Tensor | None = None, + exposure_embedding: torch.Tensor | None = None, return_output: bool = False, **unused_kwargs, ) -> torch.Tensor | DeepHealthOutput: @@ -240,16 +121,24 @@ class DeepHealth(nn.Module): event_len = event_seq.size(1) h_disease = self.token_embedding(event_seq) - h_exposure = self._encode_event_exposure( - exposure_daily=exposure_daily, - exposure_monthly=exposure_monthly, - exposure_daily_mask=exposure_daily_mask, - exposure_monthly_mask=exposure_monthly_mask, - event_shape=(event_seq.size(0), event_len), - ) - if h_exposure is not None: - h_exposure = h_exposure.to(device=event_seq.device, dtype=h_disease.dtype) - h_disease = h_disease + h_exposure + if self.use_exposure_embeddings: + if exposure_embedding is None: + raise ValueError( + "exposure_embedding is required when " + "use_exposure_embeddings=True" + ) + if exposure_embedding.shape != h_disease.shape: + raise ValueError( + "exposure_embedding must have shape " + f"{tuple(h_disease.shape)}, got {tuple(exposure_embedding.shape)}" + ) + h_disease = h_disease + exposure_embedding.to( + device=h_disease.device, dtype=h_disease.dtype + ) + elif exposure_embedding is not None: + raise ValueError( + "exposure_embedding provided but use_exposure_embeddings=False" + ) t_disease = time_seq h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype) diff --git a/train_next_step.py b/train_next_step.py index cdb95da..68d00a0 100644 --- a/train_next_step.py +++ b/train_next_step.py @@ -58,8 +58,7 @@ MODEL_INPUT_KEYS = ( ) EXPOSURE_INPUT_KEYS = ( - "exposure_daily", - "exposure_monthly", + "exposure_embedding", ) @@ -174,8 +173,7 @@ class NextStepTrainingModel(nn.Module): sex: torch.Tensor, padding_mask: torch.Tensor, readout_mask: torch.Tensor | None = None, - exposure_daily: torch.Tensor | None = None, - exposure_monthly: torch.Tensor | None = None, + exposure_embedding: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: hidden = self.model( event_seq=event_seq, @@ -183,8 +181,7 @@ class NextStepTrainingModel(nn.Module): sex=sex, padding_mask=padding_mask, target_mode="next_token", - exposure_daily=exposure_daily, - exposure_monthly=exposure_monthly, + exposure_embedding=exposure_embedding, ) if not isinstance(hidden, torch.Tensor): raise TypeError("DeepHealth forward must return a hidden-state tensor") @@ -224,19 +221,15 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--n_hist_layer", type=int, default=12) parser.add_argument("--dropout", type=float, default=0.0) parser.add_argument("--exposure_cache_dir", type=str, default=None) - parser.add_argument("--mask_onset_exposure", action="store_true") parser.add_argument( - "--d_model", - type=int, - default=64, - help="Internal TimesNet channel dimension for exposure encoding.", + "--exposure_embeddings_file", + type=str, + default=None, + help=( + "Precomputed exposure embeddings. Defaults to " + "/exposure_embeddings.npy." + ), ) - parser.add_argument("--exposure_n_layers", type=int, default=2) - parser.add_argument("--exposure_top_k", type=int, default=2) - parser.add_argument("--exposure_n_backbone_blocks", type=int, default=1) - parser.add_argument("--exposure_backbone_kernel_size", type=int, default=5) - parser.add_argument("--exposure_backbone_expansion", type=float, default=2.0) - parser.add_argument("--no_exposure_gate", action="store_true") parser.add_argument("--target_mode", type=str, default="uts", choices=["delphi2m", "uts"]) parser.add_argument("--readout_name", type=str, default=None, @@ -311,8 +304,10 @@ def parse_args() -> argparse.Namespace: raise ValueError("prefetch_factor must be positive when num_workers > 0") if args.exposure_locality_buffer_size < 0: raise ValueError("exposure_locality_buffer_size must be non-negative") - if args.d_model <= 0: - raise ValueError("--d_model must be positive") + if args.exposure_embeddings_file and not args.exposure_cache_dir: + raise ValueError( + "--exposure_cache_dir is required with --exposure_embeddings_file" + ) if args.target_mode == "uts": args.readout_name = args.readout_name or "same_time_group_end" args.include_no_event_in_uts_target = True @@ -440,14 +435,7 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth: target_mode="next_token", dist_mode="exponential", dropout=args.dropout, - use_exposure_encoder=args.exposure_cache_dir is not None, - exposure_d_model=args.d_model, - exposure_n_layers=args.exposure_n_layers, - exposure_top_k=args.exposure_top_k, - exposure_n_backbone_blocks=args.exposure_n_backbone_blocks, - exposure_backbone_kernel_size=args.exposure_backbone_kernel_size, - exposure_backbone_expansion=args.exposure_backbone_expansion, - exposure_use_gate=not args.no_exposure_gate, + use_exposure_embeddings=args.exposure_embeddings_file is not None, ) @@ -499,9 +487,8 @@ def compute_next_step_loss( "padding_mask": batch["padding_mask"], "readout_mask": batch["readout_mask"], } - if "exposure_daily" in batch: - model_kwargs["exposure_daily"] = batch["exposure_daily"] - model_kwargs["exposure_monthly"] = batch["exposure_monthly"] + if "exposure_embedding" in batch: + model_kwargs["exposure_embedding"] = batch["exposure_embedding"] logits, current_times, output_readout_mask = model(**model_kwargs) non_blocking = device.type == "cuda" targets = { @@ -631,16 +618,9 @@ def build_metadata( "dataset_metadata": { "vocab_size": int(dataset.vocab_size), }, - "use_exposure_encoder": args.exposure_cache_dir is not None, + "use_exposure_embeddings": args.exposure_embeddings_file is not None, "exposure_cache_dir": args.exposure_cache_dir, - "mask_onset_exposure": bool(args.mask_onset_exposure), - "d_model": int(args.d_model), - "exposure_n_layers": int(args.exposure_n_layers), - "exposure_top_k": int(args.exposure_top_k), - "exposure_n_backbone_blocks": int(args.exposure_n_backbone_blocks), - "exposure_backbone_kernel_size": int(args.exposure_backbone_kernel_size), - "exposure_backbone_expansion": float(args.exposure_backbone_expansion), - "exposure_use_gate": not bool(args.no_exposure_gate), + "exposure_embeddings_file": args.exposure_embeddings_file, "num_workers": int(args.num_workers), "prefetch_factor": int(args.prefetch_factor), "exposure_locality_buffer_size": int(args.exposure_locality_buffer_size), @@ -659,6 +639,10 @@ def build_metadata( def main() -> None: args = parse_args() + if args.exposure_cache_dir and args.exposure_embeddings_file is None: + args.exposure_embeddings_file = str( + Path(args.exposure_cache_dir) / "exposure_embeddings.npy" + ) device, rank, local_rank, world_size = init_distributed(args) set_seed(args.seed + rank) configure_torch_for_training(device) @@ -670,6 +654,7 @@ def main() -> None: logger.info(f"Device: {device}") logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}") logger.info(f"exposure_cache_dir={args.exposure_cache_dir}") + logger.info(f"exposure_embeddings_file={args.exposure_embeddings_file}") logger.info( "DataLoader IO: " f"num_workers={args.num_workers}, " @@ -683,8 +668,15 @@ def main() -> None: no_event_interval_years=args.no_event_interval_years, include_no_event_in_uts_target=args.include_no_event_in_uts_target, exposure_cache_dir=args.exposure_cache_dir, - mask_onset_exposure=args.mask_onset_exposure, + exposure_embeddings_file=args.exposure_embeddings_file, ) + if dataset.exposure_cache is not None: + embedding_dim = int(dataset.exposure_cache.embeddings.shape[1]) + if embedding_dim != args.n_embd: + raise ValueError( + f"Exposure embedding dim {embedding_dim} must equal " + f"--n_embd={args.n_embd}" + ) if args.train_eid_file and args.val_eid_file and args.test_eid_file: train_subset, val_subset, test_subset = split_dataset_by_eid_files( dataset=dataset,