""" 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 ( DISEASE_HISTORY_MODES, DISEASE_HISTORY_MODE_TIMED, AllFutureHealthDataset, all_future_collate_fn, ) from losses import build_loss from model_architectures import ( DEFAULT_MODEL_ARCHITECTURE, SUPPORTED_MODEL_ARCHITECTURES, ) from models import DeepHealth from targets import PAD_IDX, RESERVED_IDX from train_util import ( ContinuousRobustScalerStats, configure_torch_for_training, create_unique_run_dir, format_extra_info_types, fit_continuous_robust_scaler, get_lr, get_model_parameter_counts, load_extra_info_types_file, move_batch_to_device, resolve_device, save_checkpoint, save_config, set_optimizer_lr, set_seed, setup_logging, split_all_future_datasets, split_all_future_datasets_by_eid_files, ) MODEL_INPUT_KEYS = ( "event_seq", "time_seq", "sex", "padding_mask", "t_query", "other_type", "other_value", "other_value_kind", "other_time", ) 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("--runs_root", type=str, default="runs") parser.add_argument("--seed", type=int, default=42) parser.add_argument("--extra_info_types_file", type=str, default=None) parser.add_argument( "--disease_history_mode", type=str, default=DISEASE_HISTORY_MODE_TIMED, choices=DISEASE_HISTORY_MODES, help=( "timed=real disease times; ordered=chronological disease order with " "ordinal positions; set=unordered disease set with no disease time" ), ) 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("--train_eid_file", type=str, default="ukb_train_eid.csv") parser.add_argument("--val_eid_file", type=str, default="ukb_val_eid.csv") parser.add_argument("--test_eid_file", type=str, default="ukb_test_eid.csv") 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_layer", type=int, default=12) 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"]) parser.add_argument("--dropout", type=float, default=0.0) parser.add_argument( "--model_architecture", type=str, default=DEFAULT_MODEL_ARCHITECTURE, choices=SUPPORTED_MODEL_ARCHITECTURES, ) 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") parser.add_argument("--progress_interval", type=int, default=20) 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") use_eid_split = all( getattr(args, name) for name in ("train_eid_file", "val_eid_file", "test_eid_file") ) if not use_eid_split and 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 ) if args.disease_history_mode != DISEASE_HISTORY_MODE_TIMED: expected = { "model_architecture": "traj_mixer_v5", "time_mode": "relative", "dist_mode": "weibull", } mismatches = [ f"{name}={getattr(args, name)!r} (expected {value!r})" for name, value in expected.items() if getattr(args, name) != value ] if args.extra_info_types != []: mismatches.append( "extra_info_types must be [] via extra_info_types_none.txt" ) if mismatches: raise ValueError( f"disease_history_mode={args.disease_history_mode!r} is reserved " "for the no-extra TrajMixer + all_future + relative + Weibull " "ablation; " + "; ".join(mismatches) ) return args def build_model( args: argparse.Namespace, dataset: AllFutureHealthDataset, scaler_stats: ContinuousRobustScalerStats, ) -> DeepHealth: if tuple(int(x) for x in dataset.cont_type_ids) != scaler_stats.cont_type_ids: raise ValueError( "RobustScale statistics are not aligned with dataset.cont_type_ids" ) center = scaler_stats.center if dataset.n_cont_types > 0 else None scale = scaler_stats.scale if dataset.n_cont_types > 0 else None return DeepHealth( vocab_size=dataset.vocab_size, n_embd=args.n_embd, n_head=args.n_head, n_layer=args.n_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, continuous_value_center=center, continuous_value_scale=scale, extra_pool_reduce=args.extra_pool_reduce, target_mode="all_future", time_mode=args.time_mode, dist_mode=args.dist_mode, dropout=args.dropout, model_architecture=args.model_architecture, ) def build_criterion(args: argparse.Namespace): ignored_idx = {PAD_IDX, RESERVED_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) 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: required_keys = set(MODEL_INPUT_KEYS) required_keys.update(("future_targets", "exposure")) if args.dist_mode == "weibull": required_keys.add("future_dt") batch = move_batch_to_device( {key: batch[key] for key in required_keys}, 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"], ) 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: raise ValueError(f"Unknown dist_mode: {args.dist_mode}") 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 = torch.zeros((), device=device) n_batches = 0 skipped = 0 desc = "train" if is_train else "val" progress_interval = max(1, int(args.progress_interval)) 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 = total + loss.detach() n_batches += 1 if (batch_idx + 1) % progress_interval == 0: avg = total / max(1, n_batches) progress.set_postfix( loss=f"{float(loss.detach().cpu()):.4f}", avg=f"{float(avg.detach().cpu()):.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 float((total / max(1, n_batches)).detach().cpu()) if n_batches else float("inf") def build_metadata( args: argparse.Namespace, dataset: AllFutureHealthDataset, run_name: str, train_subset, val_subset, test_subset, scaler_stats: ContinuousRobustScalerStats, ) -> Dict[str, Any]: scaler_metadata = scaler_stats.as_metadata() return { "run_name": run_name, "dataset_class": "AllFutureHealthDataset", "collate_fn": "all_future_collate_fn", "model_class": "DeepHealth", "model_architecture": args.model_architecture, "model_target_mode": "all_future", "target_mode": "all_future", "event_stream_version": "disease_death_only_v1", "uses_assessment_event_token": False, "dist_mode": args.dist_mode, "disease_history_mode": args.disease_history_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], "continuous_value_scaling": "robust", "continuous_value_scaler": scaler_metadata, "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], "event_stream_version": "disease_death_only_v1", "uses_assessment_event_token": False, }, "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: ( ( "" if args.disease_history_mode == DISEASE_HISTORY_MODE_TIMED else f"{args.disease_history_mode}_" ) + f"{args.time_mode}_{args.dist_mode}_" f"all_future_pure_disease_{timestamp}" ), runs_root=Path(args.runs_root) / args.model_architecture, ) logger = setup_logging(run_dir) logger.info(f"Starting all-future training run: {run_name}") logger.info(f"Device: {device}") logger.info(f"Model architecture: {args.model_architecture}") logger.info(f"Disease history mode: {args.disease_history_mode}") logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}") logger.info("Continuous value scaling: RobustScale (required)") 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, disease_history_mode=args.disease_history_mode, ) 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, disease_history_mode=args.disease_history_mode, ) 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, disease_history_mode=args.disease_history_mode, ) if args.train_eid_file and args.val_eid_file and args.test_eid_file: train_subset, val_subset, test_subset = split_all_future_datasets_by_eid_files( train_dataset=train_dataset, val_dataset=val_dataset, test_dataset=test_dataset, train_eid_file=args.train_eid_file, val_eid_file=args.val_eid_file, test_eid_file=args.test_eid_file, ) logger.info( "Using eid split files: " f"train={args.train_eid_file}, val={args.val_eid_file}, test={args.test_eid_file}" ) else: 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"Using random ratio split: train={args.train_ratio}, " f"val={args.val_ratio}, test={args.test_ratio}, seed={args.seed}" ) logger.info( f"Patients/queries: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}" ) if train_dataset.n_cont_types > 0: logger.info( "Fitting continuous RobustScaler on the complete training subset: " f"patients={len(train_subset):,}, features={train_dataset.n_cont_types}" ) scaler_stats = fit_continuous_robust_scaler( train_dataset, train_subset, ) if train_dataset.n_cont_types > 0: logger.info( "Continuous RobustScaler fitted: " f"observations={int(scaler_stats.observation_count.sum()):,}, " f"min_per_feature={int(scaler_stats.observation_count.min()):,}, " f"max_per_feature={int(scaler_stats.observation_count.max()):,}" ) 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, scaler_stats=scaler_stats).to(device) parameter_counts = get_model_parameter_counts(model) logger.info( "Model parameters: " f"total={parameter_counts['model_parameter_count']:,}, " f"trainable={parameter_counts['trainable_parameter_count']:,}" ) optimizer = AdamW( model.parameters(), lr=args.base_lr, betas=tuple(args.betas), weight_decay=args.weight_decay, ) criterion = build_criterion(args) adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128) train_metadata = build_metadata( args, train_dataset, run_name, train_subset, val_subset, test_subset, scaler_stats, ) train_metadata.update(parameter_counts) save_config( args, run_dir / "train_config.json", extra=train_metadata, ) 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()