diff --git a/README.md b/README.md index 113881e..261f617 100644 --- a/README.md +++ b/README.md @@ -58,3 +58,15 @@ disease event token + pre-onset exposure embedding -> same next-token Transforme ``` The key constraint is that a disease event's own pre-onset exposure must not be used to predict that same disease event. + +Pretrain the exposure encoder as a denoising autoencoder using training-set EIDs: + +```bash +python train_exposure_autoencoder.py \ + --exposure_cache_dir ukb_exposure_cache \ + --train_eid_file ukb_train_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. diff --git a/backbones.py b/backbones.py index 678e834..9ec6be1 100644 --- a/backbones.py +++ b/backbones.py @@ -474,3 +474,117 @@ class TimesNetExposureEncoder(nn.Module): if self.gate is not None: h = torch.sigmoid(self.gate) * h return h + + +class TimesNetSequenceDecoder(nn.Module): + """Decode a fixed-size latent vector into a multivariate time series.""" + + def __init__( + self, + output_dim: int, + latent_dim: int, + d_model: int, + n_layers: int = 2, + top_k: int = 3, + n_convnext_blocks: int = 2, + conv_kernel_size: int = 7, + mlp_ratio: float = 4.0, + dropout: float = 0.0, + ): + super().__init__() + self.latent_proj = nn.Linear(latent_dim, d_model) + self.position_proj = nn.Linear(3, d_model) + self.blocks = nn.ModuleList([ + TimesNetBlock( + d_model=d_model, + top_k=top_k, + n_convnext_blocks=n_convnext_blocks, + conv_kernel_size=conv_kernel_size, + mlp_ratio=mlp_ratio, + dropout=dropout, + ) + for _ in range(n_layers) + ]) + self.final_ln = nn.LayerNorm(d_model) + self.output_proj = nn.Linear(d_model, output_dim) + + def forward(self, latent: torch.Tensor, length: int) -> torch.Tensor: + if latent.dim() != 2: + raise ValueError( + f"latent must have shape (B, D), got {tuple(latent.shape)}" + ) + position = torch.linspace( + 0.0, 1.0, length, device=latent.device, dtype=latent.dtype + ) + position = torch.stack( + [position, torch.sin(2 * torch.pi * position), + torch.cos(2 * torch.pi * position)], + dim=-1, + ) + h = self.latent_proj(latent).unsqueeze(1) + h = h + self.position_proj(position).unsqueeze(0) + for block in self.blocks: + h = block(h) + return self.output_proj(self.final_ln(h)) + + +class TimesNetExposureAutoencoder(nn.Module): + """Dual-resolution exposure autoencoder with a reusable event encoder.""" + + def __init__( + self, + n_embd: int = 120, + daily_input_dim: int = 4, + monthly_input_dim: int = 2, + d_model: int | None = None, + n_layers: int = 2, + top_k: int = 3, + n_convnext_blocks: int = 2, + conv_kernel_size: int = 7, + mlp_ratio: float = 4.0, + dropout: float = 0.0, + ): + super().__init__() + d_model = n_embd if d_model is None else d_model + encoder_kwargs = dict( + n_embd=n_embd, + daily_input_dim=daily_input_dim, + monthly_input_dim=monthly_input_dim, + d_model=d_model, + n_layers=n_layers, + top_k=top_k, + n_convnext_blocks=n_convnext_blocks, + conv_kernel_size=conv_kernel_size, + mlp_ratio=mlp_ratio, + dropout=dropout, + use_gate=True, + ) + decoder_kwargs = dict( + latent_dim=n_embd, + d_model=d_model, + n_layers=n_layers, + top_k=top_k, + n_convnext_blocks=n_convnext_blocks, + conv_kernel_size=conv_kernel_size, + mlp_ratio=mlp_ratio, + dropout=dropout, + ) + self.encoder = TimesNetExposureEncoder(**encoder_kwargs) + self.daily_decoder = TimesNetSequenceDecoder( + output_dim=daily_input_dim, **decoder_kwargs + ) + self.monthly_decoder = TimesNetSequenceDecoder( + output_dim=monthly_input_dim, **decoder_kwargs + ) + + def forward( + self, + daily: torch.Tensor, + monthly: torch.Tensor, + daily_mask: torch.Tensor | None = None, + monthly_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + latent = self.encoder(daily, monthly, daily_mask, monthly_mask) + daily_reconstruction = self.daily_decoder(latent, daily.size(1)) + monthly_reconstruction = self.monthly_decoder(latent, monthly.size(1)) + return daily_reconstruction, monthly_reconstruction, latent diff --git a/dataset.py b/dataset.py index c0d3e66..ac9d3ca 100644 --- a/dataset.py +++ b/dataset.py @@ -42,7 +42,7 @@ def _monthly_exposure_columns() -> list[str]: def _load_readonly_npy(path: Path) -> np.ndarray: - arr = np.load(path) + arr = np.load(path, mmap_mode="r") arr.setflags(write=False) return arr diff --git a/train_exposure_autoencoder.py b/train_exposure_autoencoder.py new file mode 100644 index 0000000..5f7418e --- /dev/null +++ b/train_exposure_autoencoder.py @@ -0,0 +1,330 @@ +"""Pretrain a TimesNet + ConvNeXtV2 autoencoder on training-set exposure.""" +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +import numpy as np +import torch +from torch.optim import AdamW +from torch.utils.data import DataLoader, Dataset +from tqdm import tqdm + +from backbones import TimesNetExposureAutoencoder +from dataset import ExposureCache +from train_util import ( + configure_torch_for_training, + create_unique_run_dir, + load_eid_file, + resolve_device, + set_seed, + setup_logging, +) + + +class ExposureWindowDataset(Dataset): + def __init__(self, cache: ExposureCache, row_indices: np.ndarray): + self.cache = cache + self.row_indices = np.asarray(row_indices, dtype=np.int64) + + def __len__(self) -> int: + return len(self.row_indices) + + def __getitem__(self, index: int) -> dict[str, torch.Tensor]: + row = int(self.row_indices[index]) + return { + "daily": torch.from_numpy( + np.array(self.cache.daily[row], dtype=np.float32, copy=True) + ), + "monthly": torch.from_numpy( + np.array(self.cache.monthly[row], dtype=np.float32, copy=True) + ), + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Pretrain a TimesNet + ConvNeXtV2 exposure autoencoder" + ) + parser.add_argument("--exposure_cache_dir", required=True) + parser.add_argument("--train_eid_file", default="ukb_train_eid.csv") + parser.add_argument("--runs_root", default="runs") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--val_fraction", type=float, default=0.05) + parser.add_argument("--n_embd", type=int, default=120) + parser.add_argument("--d_model", type=int, default=None) + parser.add_argument("--n_layers", type=int, default=2) + parser.add_argument("--top_k", type=int, default=3) + parser.add_argument("--n_convnext_blocks", type=int, default=2) + parser.add_argument("--conv_kernel_size", type=int, default=7) + parser.add_argument("--mlp_ratio", type=float, default=4.0) + parser.add_argument("--dropout", type=float, default=0.0) + parser.add_argument("--mask_ratio", type=float, default=0.25) + parser.add_argument("--batch_size", type=int, default=16) + parser.add_argument("--base_lr", type=float, default=3e-4) + parser.add_argument("--weight_decay", type=float, default=0.05) + parser.add_argument("--max_epochs", type=int, default=100) + parser.add_argument("--warmup_epochs", type=int, default=5) + parser.add_argument("--patience", type=int, default=12) + parser.add_argument("--grad_clip", type=float, default=1.0) + parser.add_argument("--num_workers", type=int, default=4) + parser.add_argument("--device", default="cuda") + parser.add_argument("--amp", action=argparse.BooleanOptionalAction, default=True) + args = parser.parse_args() + if not 0.0 < args.val_fraction < 1.0: + parser.error("--val_fraction must be between 0 and 1") + if not 0.0 <= args.mask_ratio < 1.0: + parser.error("--mask_ratio must be in [0, 1)") + return args + + +def select_rows( + cache: ExposureCache, train_eids: set[int], val_fraction: float, seed: int +) -> tuple[np.ndarray, np.ndarray]: + candidate_eids = np.asarray( + sorted(set(map(int, cache.eids)) & train_eids), dtype=np.int64 + ) + if len(candidate_eids) < 2: + raise ValueError("Need at least two training EIDs with cached exposure") + rng = np.random.default_rng(seed) + rng.shuffle(candidate_eids) + n_val = max(1, int(round(len(candidate_eids) * val_fraction))) + val_eids = candidate_eids[:n_val] + fit_eids = candidate_eids[n_val:] + valid_row = np.asarray(cache.row_index, dtype=np.int64) >= 0 + fit_event_rows = valid_row & np.isin(cache.eids, fit_eids) + val_event_rows = valid_row & np.isin(cache.eids, val_eids) + fit_rows = np.unique(np.asarray(cache.row_index[fit_event_rows], dtype=np.int64)) + val_rows = np.unique(np.asarray(cache.row_index[val_event_rows], dtype=np.int64)) + if len(fit_rows) == 0 or len(val_rows) == 0: + raise ValueError("Training/validation exposure rows are empty after filtering") + return fit_rows, val_rows + + +def channel_stats( + cache: ExposureCache, rows: np.ndarray, chunk_size: int = 256 +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + results = [] + for source in (cache.daily, cache.monthly): + sums = np.zeros(source.shape[-1], dtype=np.float64) + squares = np.zeros_like(sums) + counts = np.zeros_like(sums) + for start in tqdm(range(0, len(rows), chunk_size), desc="Channel statistics"): + values = np.asarray(source[rows[start:start + chunk_size]], dtype=np.float64) + finite = np.isfinite(values) + clean = np.where(finite, values, 0.0) + sums += clean.sum(axis=(0, 1)) + squares += np.square(clean).sum(axis=(0, 1)) + counts += finite.sum(axis=(0, 1)) + mean = sums / np.maximum(counts, 1.0) + variance = squares / np.maximum(counts, 1.0) - np.square(mean) + std = np.sqrt(np.maximum(variance, 1e-12)) + results.extend([mean.astype(np.float32), std.astype(np.float32)]) + return tuple(results) + + +def masked_mse( + prediction: torch.Tensor, target: torch.Tensor, mask: torch.Tensor +) -> torch.Tensor: + error = (prediction - target).square() * mask + return error.sum() / mask.sum().clamp_min(1.0) + + +def run_epoch( + model: TimesNetExposureAutoencoder, + loader: DataLoader, + device: torch.device, + stats: tuple[torch.Tensor, ...], + mask_ratio: float, + optimizer: AdamW | None, + scaler: torch.amp.GradScaler, + grad_clip: float, + amp_enabled: bool, +) -> float: + training = optimizer is not None + model.train(training) + total_loss = 0.0 + total_samples = 0 + daily_mean, daily_std, monthly_mean, monthly_std = stats + context = torch.enable_grad if training else torch.no_grad + with context(): + for batch in tqdm(loader, desc="train" if training else "val", leave=False): + daily = batch["daily"].to(device, non_blocking=True) + monthly = batch["monthly"].to(device, non_blocking=True) + daily_observed = torch.isfinite(daily) + monthly_observed = torch.isfinite(monthly) + daily = (torch.nan_to_num(daily) - daily_mean) / daily_std + monthly = (torch.nan_to_num(monthly) - monthly_mean) / monthly_std + daily = daily * daily_observed + monthly = monthly * monthly_observed + + if training and mask_ratio > 0: + daily_input_mask = daily_observed & ( + torch.rand_like(daily) >= mask_ratio + ) + monthly_input_mask = monthly_observed & ( + torch.rand_like(monthly) >= mask_ratio + ) + else: + daily_input_mask = daily_observed + monthly_input_mask = monthly_observed + daily_input = daily * daily_input_mask + monthly_input = monthly * monthly_input_mask + + if training: + optimizer.zero_grad(set_to_none=True) + with torch.autocast( + device_type=device.type, dtype=torch.float16, + enabled=amp_enabled, + ): + daily_hat, monthly_hat, _ = model( + daily_input, monthly_input, + daily_input_mask, monthly_input_mask, + ) + loss = ( + masked_mse(daily_hat, daily, daily_observed) + + masked_mse(monthly_hat, monthly, monthly_observed) + ) + if training: + scaler.scale(loss).backward() + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip) + scaler.step(optimizer) + scaler.update() + batch_size = daily.size(0) + total_loss += float(loss.detach()) * batch_size + total_samples += batch_size + return total_loss / max(total_samples, 1) + + +def learning_rate(epoch: int, args: argparse.Namespace) -> float: + if epoch < args.warmup_epochs: + return args.base_lr * (epoch + 1) / max(args.warmup_epochs, 1) + progress = (epoch - args.warmup_epochs) / max( + args.max_epochs - args.warmup_epochs - 1, 1 + ) + return args.base_lr * 0.5 * (1.0 + math.cos(math.pi * progress)) + + +def main() -> None: + args = parse_args() + set_seed(args.seed) + device = resolve_device(args.device) + configure_torch_for_training(device) + run_dir, run_name = create_unique_run_dir( + lambda stamp: f"exposure_ae_{stamp}", Path(args.runs_root) + ) + logger = setup_logging(run_dir) + cache = ExposureCache(args.exposure_cache_dir) + train_rows, val_rows = select_rows( + cache, load_eid_file(args.train_eid_file), args.val_fraction, args.seed + ) + raw_stats = channel_stats(cache, train_rows) + stats = tuple( + torch.as_tensor(value, device=device).view(1, 1, -1) + for value in raw_stats + ) + loader_kwargs = dict( + batch_size=args.batch_size, + num_workers=args.num_workers, + pin_memory=device.type == "cuda", + persistent_workers=args.num_workers > 0, + ) + train_loader = DataLoader( + ExposureWindowDataset(cache, train_rows), shuffle=True, **loader_kwargs + ) + val_loader = DataLoader( + ExposureWindowDataset(cache, val_rows), shuffle=False, **loader_kwargs + ) + model = TimesNetExposureAutoencoder( + n_embd=args.n_embd, d_model=args.d_model, n_layers=args.n_layers, + top_k=args.top_k, n_convnext_blocks=args.n_convnext_blocks, + conv_kernel_size=args.conv_kernel_size, mlp_ratio=args.mlp_ratio, + dropout=args.dropout, + ).to(device) + optimizer = AdamW( + model.parameters(), lr=args.base_lr, + weight_decay=args.weight_decay, betas=(0.9, 0.95), + ) + amp_enabled = bool(args.amp and device.type == "cuda") + scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled) + logger.info( + f"Run {run_name}: device={device}, train_rows={len(train_rows):,}, " + f"val_rows={len(val_rows):,}" + ) + config = vars(args) | { + "train_rows": len(train_rows), + "val_rows": len(val_rows), + "daily_mean": raw_stats[0].tolist(), + "daily_std": raw_stats[1].tolist(), + "monthly_mean": raw_stats[2].tolist(), + "monthly_std": raw_stats[3].tolist(), + } + (run_dir / "train_config.json").write_text( + json.dumps(config, indent=2), encoding="utf-8" + ) + + best_loss = float("inf") + stale_epochs = 0 + history = [] + for epoch in range(args.max_epochs): + lr = learning_rate(epoch, args) + for group in optimizer.param_groups: + group["lr"] = lr + train_loss = run_epoch( + model, train_loader, device, stats, args.mask_ratio, optimizer, + scaler, args.grad_clip, amp_enabled, + ) + val_loss = run_epoch( + model, val_loader, device, stats, 0.0, None, + scaler, args.grad_clip, amp_enabled, + ) + logger.info( + f"Epoch {epoch + 1:03d} | lr={lr:.3e} | " + f"train={train_loss:.6f} | val={val_loss:.6f}" + ) + history.append( + {"epoch": epoch + 1, "lr": lr, + "train_loss": train_loss, "val_loss": val_loss} + ) + if val_loss < best_loss: + best_loss = val_loss + stale_epochs = 0 + torch.save( + { + "model_state_dict": model.state_dict(), + "encoder_state_dict": model.encoder.state_dict(), + "model_config": { + key: config[key] for key in ( + "n_embd", "d_model", "n_layers", "top_k", + "n_convnext_blocks", "conv_kernel_size", + "mlp_ratio", "dropout", + ) + }, + "normalization": { + "daily_mean": raw_stats[0], + "daily_std": raw_stats[1], + "monthly_mean": raw_stats[2], + "monthly_std": raw_stats[3], + }, + "epoch": epoch + 1, + "val_loss": val_loss, + }, + run_dir / "best.pt", + ) + else: + stale_epochs += 1 + (run_dir / "history.json").write_text( + json.dumps(history, indent=2), encoding="utf-8" + ) + if stale_epochs >= args.patience: + logger.info(f"Early stopping after {epoch + 1} epochs") + break + logger.info(f"Best validation loss: {best_loss:.6f}") + logger.info(f"Checkpoint: {run_dir / 'best.pt'}") + + +if __name__ == "__main__": + main()