""" Train DeepHealth with query-conditioned all-future supervision. Training samples are patient-level. For each patient and each __getitem__ call, AllFutureHealthDataset randomly samples a query time t_query, uses events at or before t_query as history, and uses events after t_query as the future target set. Validation/test samples are deterministic query points built from future event times, then split by patient. """ from __future__ import annotations import argparse import json import logging import math import time from pathlib import Path from typing import Any, Dict import numpy as np import torch from torch.nn.utils import clip_grad_norm_ from torch.optim import AdamW from torch.utils.data import DataLoader, RandomSampler from tqdm.auto import tqdm from dataset import AllFutureHealthDataset, all_future_collate_fn from losses import build_loss from models import DeepHealth from targets import CHECKUP_IDX, PAD_IDX from train_util import ( configure_torch_for_training, create_unique_run_dir, load_extra_info_types_file, resolve_device, save_checkpoint, save_config, set_optimizer_lr, set_seed, setup_logging, split_all_future_datasets, ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Train DeepHealth with all-future supervision") parser.add_argument("--data_prefix", type=str, default="ukb") parser.add_argument("--labels_file", type=str, default="labels.csv") parser.add_argument("--seed", type=int, default=42) parser.add_argument("--extra_info_types_file", type=str, default=None) parser.add_argument("--train_ratio", type=float, default=0.7) parser.add_argument("--val_ratio", type=float, default=0.15) parser.add_argument("--test_ratio", type=float, default=0.15) parser.add_argument("--min_history_events", type=int, default=1) parser.add_argument("--min_future_events", type=int, default=1) parser.add_argument("--validation_query_seed", type=int, default=None) parser.add_argument("--n_embd", type=int, default=120) parser.add_argument("--n_head", type=int, default=10) parser.add_argument("--n_hist_layer", type=int, default=12) parser.add_argument("--n_tab_layer", type=int, default=4) parser.add_argument("--n_bins", type=int, default=16) parser.add_argument("--extra_pool_reduce", type=str, default="mean", choices=["mean", "sum"]) parser.add_argument("--time_mode", type=str, default="relative", choices=["relative", "absolute"]) parser.add_argument("--dist_mode", type=str, default="exponential", choices=["exponential", "weibull", "mixed"]) parser.add_argument("--dropout", type=float, default=0.0) parser.add_argument("--batch_size", type=int, default=128) parser.add_argument("--base_lr", type=float, default=3e-4) parser.add_argument("--weight_decay", type=float, default=0.1) parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99)) parser.add_argument("--grad_clip", type=float, default=1.0) parser.add_argument("--max_epochs", type=int, default=200) parser.add_argument("--warmup_epochs", type=int, default=10) parser.add_argument("--patience", type=int, default=15) parser.add_argument("--min_lr_ratio", type=float, default=0.1) parser.add_argument("--num_workers", type=int, default=4) parser.add_argument("--device", type=str, default="cuda") args = parser.parse_args() if args.min_history_events < 1: raise ValueError("min_history_events must be >= 1") if args.min_future_events < 1: raise ValueError("min_future_events must be >= 1") if not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0): raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0") if args.validation_query_seed is None: args.validation_query_seed = int(args.seed) args.extra_info_types = ( load_extra_info_types_file(args.extra_info_types_file) if args.extra_info_types_file is not None else None ) return args def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float: if epoch < args.warmup_epochs: return adaptive_lr * (epoch + 1) / args.warmup_epochs progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs) cosine = 0.5 * (1 + math.cos(math.pi * progress)) return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio)) def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]: non_blocking = device.type == "cuda" return { key: value.to(device, non_blocking=non_blocking) if isinstance(value, torch.Tensor) else value for key, value in batch.items() } def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> DeepHealth: return DeepHealth( vocab_size=dataset.vocab_size, n_embd=args.n_embd, n_head=args.n_head, n_hist_layer=args.n_hist_layer, n_tab_layer=args.n_tab_layer, n_types=dataset.n_types, n_cont_types=dataset.n_cont_types, n_categories=dataset.n_categories, cont_type_ids=dataset.cont_type_ids, n_bins=args.n_bins, extra_pool_reduce=args.extra_pool_reduce, target_mode="all_future", time_mode=args.time_mode, dist_mode=args.dist_mode, dropout=args.dropout, ) def build_criterion(args: argparse.Namespace, dataset: AllFutureHealthDataset): ignored_idx = {PAD_IDX, CHECKUP_IDX} if args.dist_mode == "exponential": return build_loss("exponential", ignored_idx=ignored_idx) if args.dist_mode == "weibull": return build_loss("weibull", ignored_idx=ignored_idx) if args.dist_mode == "mixed": return build_loss( "mixed", death_idx=dataset.vocab_size - 1, ignored_idx=ignored_idx, ) raise ValueError(f"Unknown dist_mode: {args.dist_mode}") def compute_all_future_loss( args: argparse.Namespace, model: DeepHealth, criterion, batch: Dict[str, torch.Tensor], device: torch.device, ) -> torch.Tensor: batch = move_batch_to_device(batch, device) hidden = model( event_seq=batch["event_seq"], time_seq=batch["time_seq"], sex=batch["sex"], padding_mask=batch["padding_mask"], t_query=batch["t_query"], other_type=batch["other_type"], other_value=batch["other_value"], other_value_kind=batch["other_value_kind"], other_time=batch["other_time"], target_mode="all_future", ) logits = model.calc_risk(hidden) if args.dist_mode == "exponential": loss = criterion( logits=logits, targets=batch["future_targets"], exposure=batch["exposure"], ) elif args.dist_mode == "weibull": loss = criterion( logits=logits, weibull_rho=model.calc_weibull_rho(hidden), targets=batch["future_targets"], dt=batch["future_dt"], exposure=batch["exposure"], ) else: loss = criterion( logits=logits, death_rho=model.calc_death_rho(hidden), targets=batch["future_targets"], dt=batch["future_dt"], exposure=batch["exposure"], ) if not torch.isfinite(loss): raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}") return loss def run_epoch( logger: logging.Logger, args: argparse.Namespace, model: DeepHealth, criterion, loader: DataLoader, optimizer: AdamW | None, device: torch.device, is_train: bool, ) -> float: model.train(is_train) total = 0.0 n_batches = 0 skipped = 0 desc = "train" if is_train else "val" progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True) for batch_idx, batch in enumerate(progress): try: loss = compute_all_future_loss(args, model, criterion, batch, device) if is_train: if optimizer is None: raise ValueError("optimizer is required for training") optimizer.zero_grad(set_to_none=True) loss.backward() if args.grad_clip > 0: clip_grad_norm_(model.parameters(), args.grad_clip) optimizer.step() total += float(loss.detach().cpu()) n_batches += 1 avg = total / max(1, n_batches) progress.set_postfix(loss=f"{float(loss.detach().cpu()):.4f}", avg=f"{avg:.4f}", skipped=skipped) except RuntimeError as exc: if "Loss is not finite" not in str(exc): raise skipped += 1 logger.warning(f"Batch {batch_idx} skipped: {str(exc)[:120]}") if skipped: logger.info(f"Skipped {skipped} batches due to non-finite loss") return total / max(1, n_batches) if n_batches else float("inf") def build_metadata( args: argparse.Namespace, dataset: AllFutureHealthDataset, run_name: str, train_subset, val_subset, test_subset, ) -> Dict[str, Any]: return { "run_name": run_name, "dataset_class": "AllFutureHealthDataset", "collate_fn": "all_future_collate_fn", "model_class": "DeepHealth", "model_target_mode": "all_future", "target_mode": "all_future", "dist_mode": args.dist_mode, "all_future_min_history_events": int(args.min_history_events), "all_future_min_future_events": int(args.min_future_events), "all_future_validation_query_seed": int(args.validation_query_seed), "extra_info_types_file": ( Path(args.extra_info_types_file).name if args.extra_info_types_file is not None else None ), "extra_info_types": [int(x) for x in dataset.extra_info_types], "dataset_metadata": { "vocab_size": int(dataset.vocab_size), "n_types": int(dataset.n_types), "n_cont_types": int(dataset.n_cont_types), "n_categories": int(dataset.n_categories), "cont_type_ids": [int(x) for x in dataset.cont_type_ids], "extra_info_types": [int(x) for x in dataset.extra_info_types], }, "split_sizes": { "train": int(len(train_subset)), "val": int(len(val_subset)), "test": int(len(test_subset)), }, "resolved_readout_name": "none", "resolved_loss_name": args.dist_mode, } 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 timestamp: f"{args.time_mode}_{args.dist_mode}_all_future_pure_disease_{timestamp}" ) logger = setup_logging(run_dir) logger.info(f"Starting all-future training run: {run_name}") logger.info(f"Device: {device}") logger.info(f"extra_info_types: {args.extra_info_types or 'all'}") logger.info("Loading all-future datasets...") train_dataset = AllFutureHealthDataset( data_prefix=args.data_prefix, labels_file=args.labels_file, split="train", min_history_events=args.min_history_events, min_future_events=args.min_future_events, validation_query_seed=args.validation_query_seed, extra_info_types=args.extra_info_types, ) val_dataset = AllFutureHealthDataset( data_prefix=args.data_prefix, labels_file=args.labels_file, split="valid", min_history_events=args.min_history_events, min_future_events=args.min_future_events, validation_query_seed=args.validation_query_seed, extra_info_types=args.extra_info_types, ) test_dataset = AllFutureHealthDataset( data_prefix=args.data_prefix, labels_file=args.labels_file, split="test", min_history_events=args.min_history_events, min_future_events=args.min_future_events, validation_query_seed=args.validation_query_seed, extra_info_types=args.extra_info_types, ) train_subset, val_subset, test_subset = split_all_future_datasets( train_dataset=train_dataset, val_dataset=val_dataset, test_dataset=test_dataset, train_ratio=args.train_ratio, val_ratio=args.val_ratio, test_ratio=args.test_ratio, seed=args.seed, ) logger.info( f"Patients/queries: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}" ) train_loader = DataLoader( train_subset, batch_size=args.batch_size, sampler=RandomSampler(train_subset, generator=torch.Generator().manual_seed(args.seed)), collate_fn=all_future_collate_fn, num_workers=args.num_workers, pin_memory=device.type == "cuda", persistent_workers=args.num_workers > 0, prefetch_factor=2 if args.num_workers > 0 else None, ) val_loader = DataLoader( val_subset, batch_size=args.batch_size, shuffle=False, collate_fn=all_future_collate_fn, num_workers=args.num_workers, pin_memory=device.type == "cuda", persistent_workers=args.num_workers > 0, prefetch_factor=2 if args.num_workers > 0 else None, ) test_loader = DataLoader( test_subset, batch_size=args.batch_size, shuffle=False, collate_fn=all_future_collate_fn, num_workers=args.num_workers, pin_memory=device.type == "cuda", persistent_workers=args.num_workers > 0, prefetch_factor=2 if args.num_workers > 0 else None, ) model = build_model(args, train_dataset).to(device) optimizer = AdamW( model.parameters(), lr=args.base_lr, betas=tuple(args.betas), weight_decay=args.weight_decay, ) criterion = build_criterion(args, train_dataset) adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128) save_config( args, run_dir / "train_config.json", extra=build_metadata(args, train_dataset, run_name, train_subset, val_subset, test_subset), ) best_val = float("inf") patience = 0 history = [] best_model_path = run_dir / "best_model.pt" start = time.time() for epoch in range(args.max_epochs): lr = get_lr(epoch, args, adaptive_lr) set_optimizer_lr(optimizer, lr) train_loss = run_epoch(logger, args, model, criterion, train_loader, optimizer, device, True) with torch.no_grad(): val_loss = run_epoch(logger, args, model, criterion, val_loader, None, device, False) is_best = val_loss < best_val if is_best: best_val = val_loss patience = 0 save_checkpoint(model, best_model_path) else: patience += 1 logger.info( f"Epoch {epoch + 1}/{args.max_epochs} | lr={lr:.6f} | " f"train_loss={train_loss:.6f} | val_loss={val_loss:.6f} | " f"best_val_loss={best_val:.6f} | patience={patience}/{args.patience} | " f"elapsed={time.time() - start:.1f}s" ) history.append({ "epoch": epoch + 1, "lr": lr, "train_loss": train_loss, "val_loss": val_loss, "best_val_loss": best_val, "is_best": int(is_best), }) if patience >= args.patience: logger.info(f"Early stopping triggered at epoch {epoch + 1}") break with (run_dir / "history.json").open("w", encoding="utf-8") as f: json.dump(history, f, indent=2) logger.info("Evaluating best model on all-future test queries...") model.load_state_dict(torch.load(best_model_path, map_location=device)) with torch.no_grad(): test_loss = run_epoch(logger, args, model, criterion, test_loader, None, device, False) logger.info(f"Test loss: {test_loss:.6f}") logger.info(f"Best checkpoint: {best_model_path}") if __name__ == "__main__": main()